Feature: Allow User to control their VM

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "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,7 @@ dependencies = [
     "pytest-cov==4.1.0",
     "pytest-mock==3.12.0",
     "pytest-asyncio==0.23.5",
+    "aioresponses==0.7.6",
     "fastapi",
     "httpx",
     "secp256k1",

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/vmclient.py

@@ -0,0 +1,188 @@
+import datetime
+import json
+import logging
+from typing import Any, Dict, Optional, Tuple
+from urllib.parse import urlparse
+
+import aiohttp
+from eth_account.messages import encode_defunct
+from jwcrypto import jwk
+from jwcrypto.jwa import JWA
+
+from aleph.sdk.types import Account
+from aleph.sdk.utils import 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 = account
+        self.ephemeral_key: jwk.JWK = jwk.JWK.generate(kty="EC", crv="P-256")
+        self.node_url: str = node_url
+        self.pubkey_payload = self._generate_pubkey_payload()
+        self.pubkey_signature_header: str = ""
+        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},
+            }
+        )
+
+    def create_payload(self, vm_id: str, operation: str) -> Dict[str, str]:
+        path = f"/control/machine/{vm_id}/{operation}"
+        payload = {
+            "time": datetime.datetime.utcnow().isoformat() + "Z",
+            "method": "POST",
+            "path": path,
+        }
+        return payload
+
+    def sign_payload(self, 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
+
+    async def _generate_header(
+        self, vm_id: str, operation: str
+    ) -> Tuple[str, Dict[str, str]]:
+        payload = self.create_payload(vm_id, operation)
+        signed_operation = self.sign_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, operation):
+        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):
+        if not self.pubkey_signature_header:
+            self.pubkey_signature_header = (
+                await self._generate_pubkey_signature_header()
+            )
+
+        payload = self.create_payload(vm_id, "logs")
+        signed_operation = self.sign_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):
+        return await self.notify_allocation(vm_id)
+
+    async def stop_instance(self, vm_id):
+        return await self.perform_operation(vm_id, "stop")
+
+    async def reboot_instance(self, vm_id):
+
+        return await self.perform_operation(vm_id, "reboot")
+
+    async def erase_instance(self, vm_id):
+        return await self.perform_operation(vm_id, "erase")
+
+    async def expire_instance(self, vm_id):
+        return await self.perform_operation(vm_id, "expire")
+
+    async def notify_allocation(self, vm_id) -> Tuple[Any, str]:
+        json_data = {"instance": vm_id}
+        async with self.session.post(
+            f"{self.node_url}/control/allocation/notify", json=json_data
+        ) as s:
+            form_response_text = await s.text()
+            return s.status, form_response_text
+
+    async def manage_instance(self, vm_id, operations):
+        for operation in operations:
+            status, response = await self.perform_operation(vm_id, operation)
+            if status != 200:
+                return status, response
+        return
+
+    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,8 @@ 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

@@ -184,3 +184,14 @@ 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

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):