Add AlephApp.vm_hash

src/aleph/sdk/vm/app.py

@@ -1,3 +1,6 @@
+import base64
+import binascii
+import socket
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -85,3 +88,29 @@ async def send_handler_result():
     def __getattr__(self, name):
         # Default all calls to the HTTP handler
         return getattr(self.http_app, name)
+
+    @property
+    def vm_hash(self) -> Optional[str]:
+        """
+        Returns the hash of the VM that is running this app. If the VM is not
+        running in Aleph, this will return None.
+        """
+        # Get hostname from environment
+        hostname = socket.gethostname()
+
+        # Add padding if necessary
+        padding_length = len(hostname) % 8
+        if padding_length != 0:
+            hostname += "=" * (8 - padding_length)
+
+        try:
+            # Convert the hostname back to its original binary form
+            item_hash_binary = base64.b32decode(hostname.upper())
+
+            # Convert the binary form to the original vm_hash
+            vm_hash = base64.b16encode(item_hash_binary).decode().lower()
+        except binascii.Error:
+            # If the hostname is not a valid base32 string, just return None
+            return None
+
+        return vm_hash

tests/unit/test_vm_app.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 
 import pytest
 from fastapi.testclient import TestClient
@@ -31,3 +32,21 @@ def test_app_http():
     response = client.get("/")
     assert response.status_code == 200
     assert response.json() == {"index": "/"}
+
+
+def test_get_vm_hash(mocker):
+    vm_hash = "deadbeef" * 8
+    # Uses the same logic as
+    # https://github.com/aleph-im/aleph-vm/blob/main/runtimes/aleph-debian-11-python/init1.py#L488
+    item_hash_binary: bytes = base64.b16decode(vm_hash.encode().upper())
+    hostname = base64.b32encode(item_hash_binary).decode().strip("=").lower()
+
+    mocker.patch("socket.gethostname", return_value=hostname)
+
+    assert app.vm_hash == vm_hash
+
+
+def test_get_vm_hash_no_vm(mocker):
+    mocker.patch("socket.gethostname", return_value="not-a-vm")
+
+    assert app.vm_hash is None