New test cases, minor refactor, more logging

Author: brunorijsman

client/client.py

@@ -11,6 +11,7 @@
 from common import utils
 from common.logging import LOGGER
 from common.user_key import UserKey
+from common.utils import key_id_str_to_uuid
 from .peer_hub import PeerHub
 
 
@@ -125,15 +126,12 @@ async def etsi_get_key(
         }
 
     async def etsi_get_key_with_key_ids(
-        self, master_sae_id: str, slave_sae_id: str, key_id: str
+        self, master_sae_id: str, slave_sae_id: str, key_id_str: str
     ):
         """
         ETSI QKD 014 V1.1.1 Get key with key IDs API.
         """
-        try:
-            key_id = UUID(key_id)
-        except ValueError as exc:
-            raise exceptions.InvalidKeyIDError(key_id) from exc
+        key_id = key_id_str_to_uuid(key_id_str)
         key = await self.gather_key_from_peer_hubs(master_sae_id, slave_sae_id, key_id)
         return {
             "keys": [

client/peer_hub.py

@@ -9,6 +9,7 @@
 from common.block import APIBlock, Block
 from common.encryption_key import EncryptionKey
 from common.logging import LOGGER
+from common.owner import Owner
 from common.pool import Pool
 from common.registration_api import (
     APIPutRegistrationRequest,
@@ -48,8 +49,8 @@ def __init__(self, client, base_url):
             self._base_url = self._base_url[:-1]
         self._registered = False
         hub_name = base_url.split("/")[-1]
-        self._local_pool = Pool(hub_name, Pool.Owner.LOCAL)
-        self._peer_pool = Pool(hub_name, Pool.Owner.PEER)
+        self._local_pool = Pool(hub_name, Owner.LOCAL)
+        self._peer_pool = Pool(hub_name, Owner.PEER)
         self._register_task = None
         self._local_pool_request_psrd_task = None
         self._peer_pool_request_psrd_task = None
@@ -165,9 +166,9 @@ async def request_psrd_task(self, pool: Pool) -> None:
             LOGGER.info(f"Finish {task_name}")
         finally:
             match pool.owner:
-                case Pool.Owner.LOCAL:
+                case Owner.LOCAL:
                     self._local_pool_request_psrd_task = None
-                case Pool.Owner.PEER:
+                case Owner.PEER:
                     self._peer_pool_request_psrd_task = None
 
     async def attempt_request_psrd(self, pool: Pool) -> bool:
@@ -177,11 +178,7 @@ async def attempt_request_psrd(self, pool: Pool) -> bool:
         """
         assert self._registered
         url = f"{self._base_url}/dske/oob/v1/psrd"
-        match pool.owner:
-            case Pool.Owner.LOCAL:
-                owner_str = "client"
-            case Pool.Owner.PEER:
-                owner_str = "hub"
+        owner_str = pool.owner.to_str(local_name="client", peer_name="hub")
         params = {
             "client_name": self._client.name,
             "owner": owner_str,

common/block.py

@@ -8,6 +8,7 @@
 import pydantic
 from bitarray import bitarray
 from common.fragment import Fragment
+from common.logging import LOGGER
 from common.utils import bytes_to_str, str_to_bytes
 from common.exceptions import (
     InvalidBlockUUIDError,
@@ -147,9 +148,14 @@ def take_data(self, start: int, size: int) -> bytes:
         Take data from the block at the specified start byte and size.
         """
         end = start + size
-        if start < 0 or end > self._size:
+        if start < 0:
+            LOGGER.error(f"Take data from blocK: invalid start index {start}")
             raise InvalidPSRDIndex(self._block_uuid, start)
+        if end > self._size:
+            LOGGER.error(f"Take data from block: invalid end index {end}")
+            raise InvalidPSRDIndex(self._block_uuid, end)
         if self._used[start:end].any():
+            LOGGER.error(f"Take data from block: already in use {start}:{end}")
             raise PSRDDataAlreadyUsedError(self._block_uuid, start, size)
         self._used[start:end] = True
         data = self._data[start:end]
@@ -185,10 +191,12 @@ def from_api(cls, api_block: APIBlock) -> "Block":
         try:
             block_uuid = UUID(api_block.block_uuid)
         except ValueError as exc:
+            LOGGER.error(f"Invalid block UUID in API block: {api_block.block_uuid}")
             raise InvalidBlockUUIDError(api_block.block_uuid) from exc
         try:
             data = str_to_bytes(api_block.data)
         except Exception as exc:
+            LOGGER.error(f"Invalid block data in API block: {api_block.data}")
             raise InvalidPSRDDataError from exc
         return Block(block_uuid, data)
 

common/exceptions.py

@@ -353,15 +353,15 @@ def __init__(self, encoded_signing_key: str):
         )
 
 
-class EncryptorNotRegisteredForClientError(DSKEException):
+class EncryptorNotConnectedToClientError(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.",
+            message="Encryptor is not connected to client.",
             details={
                 "client_name": client_name,
                 "encryptor_name": encryptor_name,
@@ -390,13 +390,13 @@ class WrongMasterSAEIDError(DSKEException):
     the master SAE ID used in the original Get key request.
     """
 
-    def __init__(self, client_name: str, master_sae_id: str, key_id: str):
+    def __init__(self, requested_master_sae_id: str, share_master_sae_id, key_id: str):
         super().__init__(
             status_code=status.HTTP_400_BAD_REQUEST,
             message="Master SAE ID does not match the one used in the original Get key request.",
             details={
-                "client_name": client_name,
-                "master_sae_id": master_sae_id,
+                "requested_master_sae_id": requested_master_sae_id,
+                "share_master_sae_id": share_master_sae_id,
                 "key_id": key_id,
             },
         )
@@ -408,13 +408,13 @@ class WrongSlaveSAEIDError(DSKEException):
     the slave SAE ID used in the original Get key request.
     """
 
-    def __init__(self, client_name: str, slave_sae_id: str, key_id: str):
+    def __init__(self, requested_master_sae_id: str, share_master_sae_id, key_id: str):
         super().__init__(
             status_code=status.HTTP_400_BAD_REQUEST,
             message="Slave SAE ID does not match the one used in the original Get key request.",
             details={
-                "client_name": client_name,
-                "slave_sae_id": slave_sae_id,
+                "requested_master_sae_id": requested_master_sae_id,
+                "share_master_sae_id": share_master_sae_id,
                 "key_id": key_id,
             },
         )

common/fragment.py

@@ -5,6 +5,7 @@
 from uuid import UUID
 import pydantic
 from common.exceptions import InvalidBlockUUIDError, InvalidEncodedFragmentError
+from common.logging import LOGGER
 from . import utils
 
 
@@ -106,6 +107,9 @@ def from_api(
         try:
             block_uuid = UUID(api_fragment.block_uuid)
         except ValueError as exc:
+            LOGGER.error(
+                "Invalid block UUID in API fragment: %s", api_fragment.block_uuid
+            )
             raise InvalidBlockUUIDError(block_uuid=api_fragment.block_uuid) from exc
         block = pool.get_block(block_uuid)
         data = block.take_data(api_fragment.start, api_fragment.size)
@@ -135,20 +139,35 @@ def from_enc_str(
         """
         parts = enc_str.split(":")
         if len(parts) != 3:
+            LOGGER.error(
+                f"Invalid encoded fragment {enc_str} "
+                f"(expected three parts separated by :)"
+            )
             raise InvalidEncodedFragmentError(encoded_fragment=enc_str)
         block_uuid_str, start_byte_str, size_str = parts
         try:
             block_uuid = UUID(block_uuid_str)
         except ValueError as exc:
+            LOGGER.error(
+                f"Invalid encoded fragment {enc_str} "
+                f"(invalid block UUID {block_uuid_str})"
+            )
             raise InvalidBlockUUIDError(block_uuid=block_uuid_str) from exc
         block = pool.get_block(block_uuid)
         try:
             start = int(start_byte_str)
         except ValueError as exc:
+            LOGGER.error(
+                f"Invalid encoded fragment {enc_str} "
+                f"(invalid start byte {start_byte_str})"
+            )
             raise InvalidEncodedFragmentError(encoded_fragment=enc_str) from exc
         try:
             size = int(size_str)
         except ValueError as exc:
+            LOGGER.error(
+                f"Invalid encoded fragment {enc_str} " f"(invalid size {size_str})"
+            )
             raise InvalidEncodedFragmentError(encoded_fragment=enc_str) from exc
         data = block.take_data(start, size)
         return Fragment(block=block, start=start, size=size, data=data)

common/owner.py

@@ -0,0 +1,44 @@
+"""
+Enumeration for pool ownership.
+"""
+
+import enum
+from .exceptions import InvalidOwnerError
+
+
+class Owner(enum.Enum):
+    """
+    Who owns a pool? The client node or the hub node? Only the owner is allowed to make
+    allocations out of the pool. The non-owner only takes data out of the pool, but the peer
+    decides which data is taken (i.e. the peer does the allocation).
+    """
+
+    LOCAL = 1
+    PEER = 2
+
+    def __str__(self):
+        return self.to_str()
+
+    def to_str(self, local_name: str = "local", peer_name: str = "peer") -> str:
+        """
+        Convert the Owner to a string.
+        """
+        match self:
+            case Owner.LOCAL:
+                return local_name
+            case Owner.PEER:
+                return peer_name
+
+    @staticmethod
+    def from_str(
+        owner_str: str, local_name: str = "local", peer_name: str = "peer"
+    ) -> "Owner":
+        """
+        Create an Owner from a string.
+        """
+        lower_owner_str = owner_str.lower()
+        if lower_owner_str == local_name:
+            return Owner.LOCAL
+        if lower_owner_str == peer_name:
+            return Owner.PEER
+        raise InvalidOwnerError(owner_str)

common/pool.py

@@ -2,33 +2,20 @@
 A pool of blocks.
 """
 
-import enum
 from uuid import UUID
 from pydantic import PositiveInt
 from .allocation import Allocation
 from .block import Block
-from .logging import LOGGER
 from .exceptions import OutOfPreSharedRandomDataError, InvalidBlockUUIDError
+from .logging import LOGGER
+from .owner import Owner
 
 
 class Pool:
     """
     A pool of blocks.
     """
 
-    class Owner(enum.Enum):
-        """
-        Who owns the pool? The client node or the hub node? Only the owner is allowed to make
-        allocations out of the pool. The non-owner only takes data out of the pool, but the peer
-        decides which data is taken (i.e. the peer does the allocation).
-        """
-
-        LOCAL = 1
-        PEER = 2
-
-        def __str__(self):
-            return self.name.lower()
-
     _name: str
     _blocks: list[Block]
     _owner: Owner
@@ -81,6 +68,7 @@ def get_block(self, block_uuid: UUID) -> Block:
         for block in self._blocks:
             if block.uuid == block_uuid:
                 return block
+        LOGGER.error(f"Block UUID not found in {str(self.owner)} pool: {block_uuid}")
         raise InvalidBlockUUIDError(block_uuid=str(block_uuid))
 
     def allocate(self, size: PositiveInt, purpose: str) -> Allocation:

common/share.py

@@ -3,6 +3,8 @@
 """
 
 from uuid import UUID
+from .exceptions import WrongMasterSAEIDError, WrongSlaveSAEIDError
+from .logging import LOGGER
 from .utils import bytes_to_str
 
 
@@ -98,3 +100,29 @@ def to_mgmt(self):
             "share_index": self._share_index,
             "value": bytes_to_str(self._value, truncate=True),
         }
+
+    def check_master_sae(self, master_sae_id: str):
+        """
+        Check if the given master SAE ID matches the stored master SAE ID.
+        Raise an exception if not.
+        """
+        if self.master_sae_id != master_sae_id:
+            key_id_str = str(self.user_key_id)
+            LOGGER.warning(
+                f"Requested master SAE ID {master_sae_id} does not match master SAE ID "
+                f"{self.master_sae_id} for key ID {key_id_str}"
+            )
+            raise WrongMasterSAEIDError(master_sae_id, self.master_sae_id, key_id_str)
+
+    def check_slave_sae(self, slave_sae_id: str):
+        """
+        Check if the given slave SAE ID matches the stored slave SAE ID.
+        Raise an exception if not.
+        """
+        if self.slave_sae_id != slave_sae_id:
+            key_id_str = str(self.user_key_id)
+            LOGGER.warning(
+                f"Requested slave SAE ID {slave_sae_id} does not match slave SAE ID "
+                f"{self.slave_sae_id} for key ID {key_id_str}"
+            )
+            raise WrongSlaveSAEIDError(slave_sae_id, self.slave_sae_id, key_id_str)

common/tests/test_pool.py

@@ -6,6 +6,7 @@
 import pytest
 from common.pool import Pool
 from common.exceptions import InvalidBlockUUIDError, OutOfPreSharedRandomDataError
+from common.owner import Owner
 from common.utils import bytes_to_str
 from .unit_test_common import create_test_pool_and_blocks
 
@@ -14,15 +15,15 @@ def test_init():
     """
     Initialize a pool.
     """
-    _pool = Pool(name="test_pool", owner=Pool.Owner.LOCAL)
+    _pool = Pool(name="test_pool", owner=Owner.LOCAL)
 
 
 def test_properties():
     """
     Properties of the pool.
     """
-    pool = Pool(name="test_pool", owner=Pool.Owner.LOCAL)
-    assert pool.owner == Pool.Owner.LOCAL
+    pool = Pool(name="test_pool", owner=Owner.LOCAL)
+    assert pool.owner == Owner.LOCAL
 
 
 def test_nr_used_and_unused_bytes():

common/tests/unit_test_common.py

@@ -5,6 +5,7 @@
 from uuid import uuid4
 from typing import List
 from common.block import Block
+from common.owner import Owner
 from common.pool import Pool
 
 
@@ -29,7 +30,7 @@ def create_test_pool_and_blocks(block_sizes: List[int]):
     """
     Create a test pool with blocks of the given sizes.
     """
-    pool = Pool(name="test_pool", owner=Pool.Owner.LOCAL)
+    pool = Pool(name="test_pool", owner=Owner.LOCAL)
     blocks = []
     for block_size in block_sizes:
         block = create_test_block(block_size)