mypy

Author: olethanh

src/aleph/sdk/client/vm_client.py

@@ -44,7 +44,7 @@ def _generate_pubkey_payload(self) -> Dict[str, Any]:
         return {
             "pubkey": json.loads(self.ephemeral_key.export_public()),
             "alg": "ECDSA",
-            "domain": urlparse(self.node_url).hostname,
+            "domain": self.node_domain,
             "address": self.account.get_address(),
             "expires": (
                 datetime.datetime.utcnow() + datetime.timedelta(days=1)
@@ -65,15 +65,15 @@ async def _generate_pubkey_signature_header(self) -> str:
                 "sender": self.account.get_address(),
                 "payload": pubkey_payload,
                 "signature": pubkey_signature,
-                "content": {"domain": urlparse(self.node_url).hostname},
+                "content": {"domain": self.node_domain},
             }
         )
 
     async def _generate_header(
         self, vm_id: ItemHash, operation: str, method: str
     ) -> Tuple[str, Dict[str, str]]:
         payload = create_vm_control_payload(
-            vm_id, operation, domain=urlparse(self.node_url).hostname, method=method
+            vm_id, operation, domain=self.node_domain, method=method
         )
         signed_operation = sign_vm_control_payload(payload, self.ephemeral_key)
 
@@ -90,6 +90,13 @@ async def _generate_header(
         path = payload["path"]
         return f"{self.node_url}{path}", headers
 
+    @property
+    def node_domain(self) -> str:
+        domain = urlparse(self.node_url).hostname
+        if not domain:
+            raise Exception("Could not parse node domain")
+        return domain
+
     async def perform_operation(
         self, vm_id: ItemHash, operation: str, method: str = "POST"
     ) -> Tuple[Optional[int], str]:
@@ -120,7 +127,7 @@ async def get_logs(self, vm_id: ItemHash) -> AsyncGenerator[str, None]:
             )
 
         payload = create_vm_control_payload(
-            vm_id, "stream_logs", method="get", domain=urlparse(self.node_url).hostname
+            vm_id, "stream_logs", method="get", domain=self.node_domain
         )
         signed_operation = sign_vm_control_payload(payload, self.ephemeral_key)
         path = payload["path"]

tests/unit/test_vm_client.py

@@ -290,7 +290,7 @@ async def test_vm_client_generate_correct_authentication_headers():
         session=aiohttp.ClientSession(),
     )
 
-    path, headers = await vm_client._generate_header(vm_id, "reboot")
+    path, headers = await vm_client._generate_header(vm_id, "reboot", method="post")
     signed_pubkey = SignedPubKeyHeader.parse_raw(headers["X-SignedPubKey"])
     signed_operation = SignedOperation.parse_raw(headers["X-SignedOperation"])
     address = verify_signed_operation(signed_operation, signed_pubkey)