Feature: Allow User to control their VM

.github/workflows/pytest.yml

@@ -14,8 +14,11 @@ on:
 jobs:
   build:
     strategy:
+      fail-fast: false
       matrix:
-        python-version: [ "3.9", "3.10", "3.11", "3.12" ]
+        python-version: [ "3.9", "3.10", "3.11" ]
+        # An issue with secp256k1 prevents Python 3.12 from working
+        # See https://github.com/baking-bad/pytezos/issues/370
     runs-on: ubuntu-latest
 
     steps:

README.md

@@ -67,7 +67,7 @@ $ pip install -e .[all]
 You can use the test env defined for hatch to run the tests:
 
 ```shell
-$ hatch run test:run
+$ hatch run testing:run
 ```
 
 See `hatch env show` for more information about all the environments and their scripts.

pyproject.toml

@@ -23,11 +23,12 @@ classifiers = [
 ]
 dependencies = [
     "aiohttp>=3.8.3",
-    "aleph-message~=0.4.4",
+    "aleph-message>=0.4.7",
     "coincurve; python_version<\"3.11\"",
     "coincurve>=19.0.0; python_version>=\"3.11\"",
     "eth_abi>=4.0.0; python_version>=\"3.11\"",
     "eth_account>=0.4.0,<0.11.0",
+    "jwcrypto==1.5.6",
     "python-magic",
     "typer",
     "typing_extensions",
@@ -122,6 +123,8 @@ dependencies = [
     "pytest-cov==4.1.0",
     "pytest-mock==3.12.0",
     "pytest-asyncio==0.23.5",
+    "pytest-aiohttp==1.0.5",
+    "aioresponses==0.7.6",
     "fastapi",
     "httpx",
     "secp256k1",
@@ -150,13 +153,13 @@ dependencies = [
 [tool.hatch.envs.linting.scripts]
 typing = "mypy --config-file=pyproject.toml {args:} ./src/ ./tests/ ./examples/"
 style = [
-    "ruff {args:.} ./src/ ./tests/ ./examples/",
+    "ruff check {args:.} ./src/ ./tests/ ./examples/",
     "black --check --diff {args:} ./src/ ./tests/ ./examples/",
     "isort --check-only --profile black {args:} ./src/ ./tests/ ./examples/",
 ]
 fmt = [
     "black {args:} ./src/ ./tests/ ./examples/",
-    "ruff --fix {args:.} ./src/ ./tests/ ./examples/",
+    "ruff check --fix {args:.} ./src/ ./tests/ ./examples/",
     "isort --profile black {args:} ./src/ ./tests/ ./examples/",
     "style",
 ]

src/aleph/sdk/chains/common.py

@@ -170,10 +170,3 @@ def get_fallback_private_key(path: Optional[Path] = None) -> bytes:
         if not default_key_path.exists():
             default_key_path.symlink_to(path)
     return private_key
-
-
-def bytes_from_hex(hex_string: str) -> bytes:
-    if hex_string.startswith("0x"):
-        hex_string = hex_string[2:]
-    hex_string = bytes.fromhex(hex_string)
-    return hex_string

src/aleph/sdk/chains/ethereum.py

@@ -7,12 +7,8 @@
 from eth_keys.exceptions import BadSignature as EthBadSignatureError
 
 from ..exceptions import BadSignatureError
-from .common import (
-    BaseAccount,
-    bytes_from_hex,
-    get_fallback_private_key,
-    get_public_key,
-)
+from ..utils import bytes_from_hex
+from .common import BaseAccount, get_fallback_private_key, get_public_key
 
 
 class ETHAccount(BaseAccount):

src/aleph/sdk/chains/substrate.py

@@ -9,7 +9,8 @@
 
 from ..conf import settings
 from ..exceptions import BadSignatureError
-from .common import BaseAccount, bytes_from_hex, get_verification_buffer
+from ..utils import bytes_from_hex
+from .common import BaseAccount, get_verification_buffer
 
 logger = logging.getLogger(__name__)
 

src/aleph/sdk/client/authenticated_http.py

@@ -30,6 +30,7 @@
 from aleph_message.models.execution.environment import (
     FunctionEnvironment,
     HypervisorType,
+    InstanceEnvironment,
     MachineResources,
 )
 from aleph_message.models.execution.instance import RootfsVolume
@@ -534,16 +535,17 @@ async def create_instance(
         timeout_seconds = timeout_seconds or settings.DEFAULT_VM_TIMEOUT
 
         payment = payment or Payment(chain=Chain.ETH, type=PaymentType.hold)
-        hypervisor = hypervisor or HypervisorType.firecracker
+
+        # Default to the QEMU hypervisor for instances.
+        selected_hypervisor: HypervisorType = hypervisor or HypervisorType.qemu
 
         content = InstanceContent(
             address=address,
             allow_amend=allow_amend,
-            environment=FunctionEnvironment(
-                reproducible=False,
+            environment=InstanceEnvironment(
                 internet=internet,
                 aleph_api=aleph_api,
-                hypervisor=hypervisor,
+                hypervisor=selected_hypervisor,
             ),
             variables=environment_variables,
             resources=MachineResources(

src/aleph/sdk/client/vmclient.py

@@ -0,0 +1,177 @@
+import datetime
+import json
+import logging
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
+from urllib.parse import urlparse
+
+import aiohttp
+from aleph_message.models import ItemHash
+from eth_account.messages import encode_defunct
+from jwcrypto import jwk
+
+from aleph.sdk.types import Account
+from aleph.sdk.utils import (
+    create_vm_control_payload,
+    sign_vm_control_payload,
+    to_0x_hex,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class VmClient:
+    account: Account
+    ephemeral_key: jwk.JWK
+    node_url: str
+    pubkey_payload: Dict[str, Any]
+    pubkey_signature_header: str
+    session: aiohttp.ClientSession
+
+    def __init__(
+        self,
+        account: Account,
+        node_url: str = "",
+        session: Optional[aiohttp.ClientSession] = None,
+    ):
+        self.account = account
+        self.ephemeral_key = jwk.JWK.generate(kty="EC", crv="P-256")
+        self.node_url = node_url
+        self.pubkey_payload = self._generate_pubkey_payload()
+        self.pubkey_signature_header = ""
+        self.session = session or aiohttp.ClientSession()
+
+    def _generate_pubkey_payload(self) -> Dict[str, Any]:
+        return {
+            "pubkey": json.loads(self.ephemeral_key.export_public()),
+            "alg": "ECDSA",
+            "domain": urlparse(self.node_url).netloc,
+            "address": self.account.get_address(),
+            "expires": (
+                datetime.datetime.utcnow() + datetime.timedelta(days=1)
+            ).isoformat()
+            + "Z",
+        }
+
+    async def _generate_pubkey_signature_header(self) -> str:
+        pubkey_payload = json.dumps(self.pubkey_payload).encode("utf-8").hex()
+        signable_message = encode_defunct(hexstr=pubkey_payload)
+        buffer_to_sign = signable_message.body
+
+        signed_message = await self.account.sign_raw(buffer_to_sign)
+        pubkey_signature = to_0x_hex(signed_message)
+
+        return json.dumps(
+            {
+                "sender": self.account.get_address(),
+                "payload": pubkey_payload,
+                "signature": pubkey_signature,
+                "content": {"domain": urlparse(self.node_url).netloc},
+            }
+        )
+
+    async def _generate_header(
+        self, vm_id: ItemHash, operation: str
+    ) -> Tuple[str, Dict[str, str]]:
+        payload = create_vm_control_payload(vm_id, operation)
+        signed_operation = sign_vm_control_payload(payload, self.ephemeral_key)
+
+        if not self.pubkey_signature_header:
+            self.pubkey_signature_header = (
+                await self._generate_pubkey_signature_header()
+            )
+
+        headers = {
+            "X-SignedPubKey": self.pubkey_signature_header,
+            "X-SignedOperation": signed_operation,
+        }
+
+        path = payload["path"]
+        return f"{self.node_url}{path}", headers
+
+    async def perform_operation(
+        self, vm_id: ItemHash, operation: str
+    ) -> Tuple[Optional[int], str]:
+        if not self.pubkey_signature_header:
+            self.pubkey_signature_header = (
+                await self._generate_pubkey_signature_header()
+            )
+
+        url, header = await self._generate_header(vm_id=vm_id, operation=operation)
+
+        try:
+            async with self.session.post(url, headers=header) as response:
+                response_text = await response.text()
+                return response.status, response_text
+
+        except aiohttp.ClientError as e:
+            logger.error(f"HTTP error during operation {operation}: {str(e)}")
+            return None, str(e)
+
+    async def get_logs(self, vm_id: ItemHash) -> AsyncGenerator[str, None]:
+        if not self.pubkey_signature_header:
+            self.pubkey_signature_header = (
+                await self._generate_pubkey_signature_header()
+            )
+
+        payload = create_vm_control_payload(vm_id, "stream_logs")
+        signed_operation = sign_vm_control_payload(payload, self.ephemeral_key)
+        path = payload["path"]
+        ws_url = f"{self.node_url}{path}"
+
+        async with self.session.ws_connect(ws_url) as ws:
+            auth_message = {
+                "auth": {
+                    "X-SignedPubKey": self.pubkey_signature_header,
+                    "X-SignedOperation": signed_operation,
+                }
+            }
+            await ws.send_json(auth_message)
+
+            async for msg in ws:  # msg is of type aiohttp.WSMessage
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    yield msg.data
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    break
+
+    async def start_instance(self, vm_id: ItemHash) -> Tuple[int, str]:
+        return await self.notify_allocation(vm_id)
+
+    async def stop_instance(self, vm_id: ItemHash) -> Tuple[Optional[int], str]:
+        return await self.perform_operation(vm_id, "stop")
+
+    async def reboot_instance(self, vm_id: ItemHash) -> Tuple[Optional[int], str]:
+        return await self.perform_operation(vm_id, "reboot")
+
+    async def erase_instance(self, vm_id: ItemHash) -> Tuple[Optional[int], str]:
+        return await self.perform_operation(vm_id, "erase")
+
+    async def expire_instance(self, vm_id: ItemHash) -> Tuple[Optional[int], str]:
+        return await self.perform_operation(vm_id, "expire")
+
+    async def notify_allocation(self, vm_id: ItemHash) -> Tuple[int, str]:
+        json_data = {"instance": vm_id}
+
+        async with self.session.post(
+            f"{self.node_url}/control/allocation/notify", json=json_data
+        ) as session:
+            form_response_text = await session.text()
+
+            return session.status, form_response_text
+
+    async def manage_instance(
+        self, vm_id: ItemHash, operations: List[str]
+    ) -> Tuple[int, str]:
+        for operation in operations:
+            status, response = await self.perform_operation(vm_id, operation)
+            if status != 200 and status:
+                return status, response
+        return 200, "All operations completed successfully"
+
+    async def close(self):
+        await self.session.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.close()

src/aleph/sdk/types.py

@@ -20,6 +20,9 @@ class Account(Protocol):
     @abstractmethod
     async def sign_message(self, message: Dict) -> Dict: ...
 
+    @abstractmethod
+    async def sign_raw(self, buffer: bytes) -> bytes: ...
+
     @abstractmethod
     def get_address(self) -> str: ...
 

src/aleph/sdk/utils.py

@@ -1,5 +1,6 @@
 import errno
 import hashlib
+import json
 import logging
 import os
 from datetime import date, datetime, time
@@ -8,6 +9,7 @@
 from shutil import make_archive
 from typing import (
     Any,
+    Dict,
     Iterable,
     Mapping,
     Optional,
@@ -20,9 +22,10 @@
 )
 from zipfile import BadZipFile, ZipFile
 
-from aleph_message.models import MessageType
+from aleph_message.models import ItemHash, MessageType
 from aleph_message.models.execution.program import Encoding
 from aleph_message.models.execution.volume import MachineVolume
+from jwcrypto.jwa import JWA
 from pydantic.json import pydantic_encoder
 
 from aleph.sdk.conf import settings
@@ -184,3 +187,36 @@ def parse_volume(volume_dict: Union[Mapping, MachineVolume]) -> MachineVolume:
 def compute_sha256(s: str) -> str:
     """Compute the SHA256 hash of a string."""
     return hashlib.sha256(s.encode()).hexdigest()
+
+
+def to_0x_hex(b: bytes) -> str:
+    return "0x" + bytes.hex(b)
+
+
+def bytes_from_hex(hex_string: str) -> bytes:
+    if hex_string.startswith("0x"):
+        hex_string = hex_string[2:]
+    hex_string = bytes.fromhex(hex_string)
+    return hex_string
+
+
+def create_vm_control_payload(vm_id: ItemHash, operation: str) -> Dict[str, str]:
+    path = f"/control/machine/{vm_id}/{operation}"
+    payload = {
+        "time": datetime.utcnow().isoformat() + "Z",
+        "method": "POST",
+        "path": path,
+    }
+    return payload
+
+
+def sign_vm_control_payload(payload: Dict[str, str], ephemeral_key) -> str:
+    payload_as_bytes = json.dumps(payload).encode("utf-8")
+    payload_signature = JWA.signing_alg("ES256").sign(ephemeral_key, payload_as_bytes)
+    signed_operation = json.dumps(
+        {
+            "payload": payload_as_bytes.hex(),
+            "signature": payload_signature.hex(),
+        }
+    )
+    return signed_operation

src/aleph/sdk/wallets/ledger/ethereum.py

@@ -9,7 +9,8 @@
 from ledgereth.messages import sign_message
 from ledgereth.objects import LedgerAccount, SignedMessage
 
-from ...chains.common import BaseAccount, bytes_from_hex, get_verification_buffer
+from ...chains.common import BaseAccount, get_verification_buffer
+from ...utils import bytes_from_hex
 
 
 class LedgerETHAccount(BaseAccount):