diff --git a/src/aleph/sdk/vm/app.py b/src/aleph/sdk/vm/app.py index 88bf3eb8..0ca53807 100644 --- a/src/aleph/sdk/vm/app.py +++ b/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 diff --git a/tests/unit/test_vm_app.py b/tests/unit/test_vm_app.py index c9ed5aa9..80b2d156 100644 --- a/tests/unit/test_vm_app.py +++ b/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