SAE progress

Author: brunorijsman

client/__main__.py

@@ -6,11 +6,12 @@
 import contextlib
 import os
 import signal
+from typing import Annotated
 import fastapi
 import uvicorn
 from common import configuration
 from common import utils
-from common.exceptions import DSKEException
+from common.exceptions import DSKEException, EncryptorNotRegisteredForClientError
 from .client import Client
 
 
@@ -71,30 +72,42 @@ async def dske_exception_handler(_request: fastapi.Request, exc: DSKEException):
 
 
 @_APP.get(f"/client/{_CLIENT.name}/etsi/api/v1/keys/{{slave_sae_id}}/status")
-async def get_etsi_status(slave_sae_id: str):
+async def get_etsi_status(
+    slave_sae_id: str,
+    authorization_header: Annotated[str | None, fastapi.Header()] = None,
+):
     """
     ETSI QKD 014 API: Status.
     """
-    return await _CLIENT.etsi_status(slave_sae_id)
+    master_sae_id = calling_sae_id(authorization_header)
+    return await _CLIENT.etsi_status(master_sae_id, slave_sae_id)
 
 
 @_APP.get(f"/client/{_CLIENT.name}/etsi/api/v1/keys/{{slave_sae_id}}/enc_keys")
-async def get_etsi_get_key(slave_sae_id: str, size: int | None = None):
+async def get_etsi_get_key(
+    slave_sae_id: str,
+    size: int | None = None,
+    authorization_header: Annotated[str | None, fastapi.Header()] = None,
+):
     """
     ETSI QKD 014 API: Get Key.
     """
-    master_sae_id = _CLIENT.name  # Currently, SAE names are the same as client names.
+    master_sae_id = calling_sae_id(authorization_header)
     return await _CLIENT.etsi_get_key(master_sae_id, slave_sae_id, size)
 
 
 @_APP.get(f"/client/{_CLIENT.name}/etsi/api/v1/keys/{{master_sae_id}}/dec_keys")
-async def get_eti_get_key_with_key_ids(master_sae_id: str, key_ID: str):
+async def get_eti_get_key_with_key_ids(
+    master_sae_id: str,
+    key_ID: str,
+    authorization_header: Annotated[str | None, fastapi.Header()] = None,
+):
     """
     ETSI QKD 014 API: Get Key with Key IDs.
     """
     # ETSI QKD 014 says that ID in key_ID has to be upper case, which lint doesn't like.
     # pylint: disable=invalid-name
-    slave_sae_id = _CLIENT.name  # Currently, SAE names are the same as client names.
+    slave_sae_id = calling_sae_id(authorization_header)
     return await _CLIENT.etsi_get_key_with_key_ids(master_sae_id, slave_sae_id, key_ID)
 
 
@@ -116,6 +129,26 @@ async def post_mgmt_stop():
     return {"result": "Client stopped"}
 
 
+def calling_sae_id(authorization_header: str | None) -> str:
+    """
+    Get the SAE ID of the calling entity from the request headers.
+    """
+    if authorization_header is None:
+        if len(_CLIENT.encryptor_names) == 1:
+            # If the client has exactly one encryptor, we assume that this is the one.
+            # This makes it easier for users to use curl for testing in simple topologies.
+            master_sae_id = _CLIENT.encryptor_names[0]
+        else:
+            # There is more than one encryptor, so we cannot guess which one is calling.
+            raise EncryptorNotRegisteredForClientError(
+                client_name=_CLIENT.name,
+                encryptor_name=authorization_header,
+            )
+    else:
+        master_sae_id = authorization_header.strip()
+    return master_sae_id
+
+
 def main():
     """
     Main entry point for the hub package.

client/client.py

@@ -2,6 +2,7 @@
 A DSKE client, or just client for short.
 """
 
+import traceback
 import asyncio
 from uuid import UUID
 from common import exceptions
@@ -64,15 +65,15 @@ def to_mgmt(self):
             "peer_hubs": peer_hubs_status,
         }
 
-    async def etsi_status(self, slave_sae_id: str):
+    async def etsi_status(self, master_sae_id: str, slave_sae_id: str):
         """
         ETSI QKD 014 V1.1.1 Status API.
         """
         # See remark about ETSI QKD API in file TODO
         return {
             "source_kme_id": self._name,
-            "target_kme_id": slave_sae_id,
-            "master_sae_id": self._name,
+            "target_kme_id": "TODO",  # TODO: Determine slave KME ID from slave SAE ID
+            "master_sae_id": master_sae_id,
             "slave_sae_id": slave_sae_id,
             "key_size": self._default_key_size_in_bits,
             "stored_key_count": 25000,  # TODO

client/http_client.py

@@ -41,6 +41,12 @@ async def async_auth_flow(self, request):
             signature.add_to_headers(request.headers)
             response = yield request
             received_signature = Signature.from_headers(response.headers)
+            if received_signature is None:
+                # If the signature is missing on an error response, keep the original error response
+                # instead of raising an InvalidSignatureError which wipes out any useful error info.
+                if response.status_code < 400:
+                    raise InvalidSignatureError()
+                return
             allocation = Allocation.from_enc_str(
                 received_signature.signing_key_allocation_enc_str, self._peer_pool
             )
@@ -166,7 +172,12 @@ async def _put_or_post(
                     exception=str(exc),
                 ) from exc
             if response.status_code != 200:
-                LOGGER.error(f"Call {method} {url} {response.status_code}")
+                message = ""
+                try:
+                    message = " " + response.json().get("message")
+                except Exception:  # pylint: disable=broad-except
+                    pass
+                LOGGER.error(f"Call {method} {url} {response.status_code}{message}")
                 raise exceptions.HTTPError(
                     method=method,
                     url=url,

common/exceptions.py

@@ -317,3 +317,19 @@ def __init__(self, encoded_fragment: str):
             message="Invalid encoded fragment.",
             details={"encoded_fragment": encoded_fragment},
         )
+
+
+class EncryptorNotRegisteredForClientError(DSKEException):
+    """
+    Exception raised when an encryptor is not registered for a client.
+    """
+
+    def __init__(self, client_name: str, encryptor_name: str):
+        super().__init__(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            message="Encryptor is not registered for client.",
+            details={
+                "client_name": client_name,
+                "encryptor_name": encryptor_name,
+            },
+        )

common/signature.py

@@ -2,6 +2,7 @@
 The signature for a DSKE in-band protocol message.
 """
 
+from typing import Optional
 from .utils import bytes_to_str, str_to_bytes
 
 HEADER_NAME = "DSKE-Signature"
@@ -74,12 +75,13 @@ def add_to_headers(self, headers: dict[str, str]):
         headers[HEADER_NAME] = self.to_enc_str()
 
     @classmethod
-    def from_headers(cls, headers: dict[str, str]) -> "Signature":
+    def from_headers(cls, headers: dict[str, str]) -> Optional["Signature"]:
         """
-        Create a Signature from the DSKE-Signature header in a dictionary of HTTP headers.
+        Create a Signature from the DSKE-Signature header in a dictionary of HTTP headers. Returns
+        None if the header is not present.
         """
         signature_enc_str = headers.get(LOWER_HEADER_NAME, None)
         if signature_enc_str is None:
-            assert False  # TODO: Raise an exception instead
+            return None
         signature = cls.from_enc_str(signature_enc_str)
         return signature

common/signing_key.py

@@ -113,7 +113,7 @@ def extract_from_headers(cls, headers: dict[str, str]) -> "SigningKey":
         """
         signing_key_enc_str = headers.get(LOWER_HEADER_NAME, None)
         if signing_key_enc_str is None:
-            assert False  # TODO: Raise an exception instead
+            return None
         signing_key = cls.from_enc_str(signing_key_enc_str)
         del headers[LOWER_HEADER_NAME]
         return signing_key

hub/__main__.py

@@ -9,6 +9,7 @@
 from common import configuration
 from common import utils
 from common.block import APIBlock
+from common.exceptions import DSKEException
 from common.share_api import APIGetShareResponse, APIPostShareRequest
 from common.signing_key import MiddlewareSigningKey
 from common.registration_api import (
@@ -67,6 +68,11 @@ async def middleware_add_response_signature(
     Add a signature to the response.
     """
     signing_key = MiddlewareSigningKey.extract_from_headers(response.headers)
+    if signing_key is None:
+        # Don't sign if the signing key is missing. This could happen, for example, in
+        # error responses. The other side can always reject the response if it doesn't like
+        # the missing signature.
+        return response
     chunks = []
     async for chunk in response.body_iterator:
         chunks.append(chunk)
@@ -82,6 +88,18 @@ async def middleware_add_response_signature(
     return signed_response
 
 
+@_APP.exception_handler(DSKEException)
+async def dske_exception_handler(_request: fastapi.Request, exc: DSKEException):
+    """
+    Handle DSKE exceptions.
+    """
+    # Error responses are not signed.
+    return fastapi.responses.JSONResponse(
+        status_code=exc.status_code,
+        content={"message": exc.message, "details": exc.details},
+    )
+
+
 @_APP.put(f"/hub/{_HUB.name}/dske/oob/v1/registration")
 async def put_oob_client_registration(
     registration_request: APIPutRegistrationRequest,
@@ -106,8 +124,6 @@ async def get_oob_psrd(
     """
     DSKE Out of band: Get a block of Pre-Shared Random Data (PSRD).
     """
-    # TODO: Error if the client was not peer.
-    # TODO: Allow size to be None (use default size decided by hub).
     block = _HUB.generate_block_for_client(client_name, pool_owner, size)
     return block.to_api()
 

hub/hub.py

@@ -13,6 +13,8 @@
 from common.allocation import Allocation
 from common.block import Block
 from common.encryption_key import EncryptionKey
+from common.exceptions import EncryptorNotRegisteredForClientError
+from common.logging import LOGGER
 from common.pool import Pool
 from common.share import Share
 from common.share_api import APIGetShareResponse, APIPostShareRequest
@@ -75,6 +77,7 @@ def generate_block_for_client(
         Generate a block of PSRD for a peer client.
         """
         if client_name not in self._peer_clients:
+            LOGGER.warning(f"Peer client '{client_name}' not found")
             raise exceptions.ClientNotRegisteredError(client_name)
         peer_client = self._peer_clients[client_name]
         match pool_owner_str.lower():
@@ -83,6 +86,9 @@ def generate_block_for_client(
             case "hub":
                 pool_owner = Pool.Owner.LOCAL
             case _:
+                LOGGER.warning(
+                    f"Invalid pool owner {pool_owner_str} for peer client {client_name}"
+                )
                 raise exceptions.InvalidPoolOwnerError(pool_owner_str)
         block = peer_client.create_random_block(pool_owner, size)
         return block
@@ -96,11 +102,22 @@ async def store_share_received_from_client(
         """
         Store a key share posted by a client.
         """
+        # Lookup the peer client
         client_name = api_post_share_request.master_client_name
         if client_name not in self._peer_clients:
+            LOGGER.warning(f"Peer client {client_name} not found")
             raise exceptions.ClientNotRegisteredError(client_name)
         peer_client = self._peer_clients[client_name]
+        # Verify the request signature
         await peer_client.check_request_signature(raw_request)
+        # Check that the master encryptor (SAE) is one that was registered for the client
+        master_sae_id = api_post_share_request.master_sae_id
+        if master_sae_id not in peer_client.encryptor_names:
+            LOGGER.warning(
+                f"Encryptor {master_sae_id} not registered for client {client_name}"
+            )
+            raise EncryptorNotRegisteredForClientError(client_name, master_sae_id)
+        # Decrypt the share value
         encryption_key_allocation = Allocation.from_api(
             api_post_share_request.encryption_key_allocation, peer_client.peer_pool
         )
@@ -120,6 +137,7 @@ async def store_share_received_from_client(
         # TODO: Check if the key UUID is already present, and if so, do something sensible
         self._shares[share.user_key_id] = share
         peer_client.add_dske_signing_key_header_to_response(headers_temp_response)
+        # Clean up fully used blocks
         peer_client.delete_fully_used_blocks()
 
     async def get_share_requested_by_client(
@@ -132,25 +150,35 @@ async def get_share_requested_by_client(
         """
         Get a key share.
         """
+        # Lookup the peer client
+        if client_name not in self._peer_clients:
+            LOGGER.warning(f"Peer client {client_name} not found")
+            raise exceptions.ClientNotRegisteredError(client_name)
+        peer_client = self._peer_clients[client_name]
+        # Verify the request signature
+        await peer_client.check_request_signature(raw_request)
+        # Lookup the share
         try:
             key_id = UUID(key_id_str)
         except ValueError as exc:
+            LOGGER.warning(f"Invalid key ID {key_id_str}")
             raise exceptions.InvalidKeyIDError(key_id_str) from exc
-        # TODO: Error handling: share is not in the store
         try:
             share = self._shares[key_id]
         except KeyError as exc:
+            LOGGER.warning(f"No share for key ID {key_id_str}")
             raise exceptions.UnknownKeyIDError(key_id) from exc
-        peer_client = self._peer_clients[client_name]
-        await peer_client.check_request_signature(raw_request)
+        # Encrypt the share value
         encryption_key = EncryptionKey.from_pool(peer_client.local_pool, share.size)
         encrypted_share_value = encryption_key.encrypt(share.value)
+        # Prepare the response
         response = APIGetShareResponse(
             share_index=share.share_index,
             encryption_key_allocation=encryption_key.allocation.to_api(),
             encrypted_share_value=bytes_to_str(encrypted_share_value),
         )
         peer_client.add_dske_signing_key_header_to_response(headers_temp_response)
+        # Clean up fully used blocks
         peer_client.delete_fully_used_blocks()
         return response
 

hub/peer_client.py

@@ -7,6 +7,7 @@
 from common.allocation import Allocation
 from common.block import Block
 from common.exceptions import InvalidSignatureError
+from common.logging import LOGGER
 from common.pool import Pool
 from common.signature import Signature
 from common.signing_key import SigningKey
@@ -27,7 +28,20 @@ def __init__(self, client_name: str, encryptor_names: List[str]):
         self._encryptor_names = encryptor_names
         self._local_pool = Pool(client_name, Pool.Owner.LOCAL)
         self._peer_pool = Pool(client_name, Pool.Owner.PEER)
-        self._shares = {}
+
+    @property
+    def client_name(self) -> str:
+        """
+        Get the client name.
+        """
+        return self._client_name
+
+    @property
+    def encryptor_names(self) -> List[str]:
+        """
+        Get the list of encryptor names registered for this client.
+        """
+        return self._encryptor_names
 
     @property
     def local_pool(self) -> Pool:
@@ -93,6 +107,9 @@ async def check_request_signature(self, raw_request: fastapi.Request):
         signature_ok = received_signature.same_as(computed_signature)
         if not signature_ok:
             # TODO: Give allocation back to pool
+            LOGGER.warning(
+                f"Invalid signature received from peer client '{self._client_name}'"
+            )
             raise InvalidSignatureError()
 
     def delete_fully_used_blocks(self) -> None: