comments

Author: xanderbailey

pyiceberg/encryption/kms.py

@@ -14,14 +14,14 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Key management client interface for table encryption, and an in-memory implementation."""
+"""Key management client interface for table encryption."""
 
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 
-from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey
+from pyiceberg.typedef import EMPTY_DICT, Properties
 
 
 @dataclass(frozen=True)
@@ -36,8 +36,16 @@ class KeyManagementClient(ABC):
     """A base class for key management service implementations.
 
     Wraps and unwraps table encryption keys using master keys that the service holds.
+
+    Implementations are loaded by name from the catalog properties, so a subclass must keep
+    this constructor signature, as `FileIO` does.
     """
 
+    properties: Properties
+
+    def __init__(self, properties: Properties = EMPTY_DICT) -> None:
+        self.properties = properties
+
     @abstractmethod
     def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes:
         """Wrap a key using the master key identified by `wrapping_key_id`.
@@ -67,48 +75,3 @@ def generate_key(self, wrapping_key_id: str) -> GeneratedKey:
             wrapping_key_id (str): Identifies the master key held by the service.
         """
         raise NotImplementedError(f"{type(self).__name__} does not support key generation")
-
-
-class MemoryKeyManagementClient(KeyManagementClient):
-    """A key management service that holds its master keys in memory, for testing and demonstration.
-
-    Master keys live only in this process, with no durability or access control, so this is
-    not for production use. Mirrors Java's `MemoryMockKMS` and iceberg-rust's
-    `MemoryKeyManagementClient`.
-    """
-
-    def __init__(self, master_key_size: AesKeySize = AesKeySize.BITS_128) -> None:
-        self._master_key_size = master_key_size
-        self._master_keys: dict[str, SecureKey] = {}
-
-    def __repr__(self) -> str:
-        """Return a representation that counts the master keys without exposing them."""
-        return f"MemoryKeyManagementClient(master_key_size={self._master_key_size!r}, key_count={len(self._master_keys)})"
-
-    def add_master_key(self, wrapping_key_id: str, key: SecureKey | None = None) -> SecureKey:
-        """Register a master key under `wrapping_key_id`, generating one when `key` is omitted.
-
-        Args:
-            wrapping_key_id (str): The id to register the master key under.
-            key (SecureKey | None): Known key material, for tests that share it with another client.
-        """
-        if wrapping_key_id in self._master_keys:
-            raise ValueError(f"Master key already exists: {wrapping_key_id}")
-
-        master_key = SecureKey.generate(self._master_key_size) if key is None else key
-        self._master_keys[wrapping_key_id] = master_key
-        return master_key
-
-    def _cipher(self, wrapping_key_id: str) -> AesGcmCipher:
-        if (master_key := self._master_keys.get(wrapping_key_id)) is None:
-            raise ValueError(f"Master key not found: {wrapping_key_id}")
-
-        return AesGcmCipher(master_key)
-
-    def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes:
-        """Wrap a key with the registered master key, without AAD, as Java and iceberg-rust do."""
-        return self._cipher(wrapping_key_id).encrypt(key)
-
-    def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes:
-        """Unwrap a key wrapped by `wrap_key`."""
-        return self._cipher(wrapping_key_id).decrypt(wrapped_key)

tests/encryption/memory_kms.py

@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""An in-memory key management client for tests, mirroring Java's `MemoryMockKMS`."""
+
+from __future__ import annotations
+
+from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey
+from pyiceberg.encryption.kms import KeyManagementClient
+from pyiceberg.typedef import EMPTY_DICT, Properties
+
+
+class MemoryKeyManagementClient(KeyManagementClient):
+    """A key management service that holds its master keys in memory.
+
+    Master keys live only in this process, with no durability or access control, so this is
+    for tests only.
+    """
+
+    def __init__(self, properties: Properties = EMPTY_DICT, *, master_key_size: AesKeySize = AesKeySize.BITS_128) -> None:
+        super().__init__(properties)
+        self._master_key_size = master_key_size
+        self._master_keys: dict[str, SecureKey] = {}
+
+    def __repr__(self) -> str:
+        """Return a representation that counts the master keys without exposing them."""
+        return f"MemoryKeyManagementClient(master_key_size={self._master_key_size!r}, key_count={len(self._master_keys)})"
+
+    def add_master_key(self, wrapping_key_id: str, key: SecureKey | None = None) -> SecureKey:
+        """Register a master key under `wrapping_key_id`, generating one when `key` is omitted.
+
+        Args:
+            wrapping_key_id (str): The id to register the master key under.
+            key (SecureKey | None): Known key material, for tests that share it with another client.
+        """
+        if wrapping_key_id in self._master_keys:
+            raise ValueError(f"Master key already exists: {wrapping_key_id}")
+
+        master_key = SecureKey.generate(self._master_key_size) if key is None else key
+        self._master_keys[wrapping_key_id] = master_key
+        return master_key
+
+    def _cipher(self, wrapping_key_id: str) -> AesGcmCipher:
+        if (master_key := self._master_keys.get(wrapping_key_id)) is None:
+            raise ValueError(f"Master key not found: {wrapping_key_id}")
+
+        return AesGcmCipher(master_key)
+
+    def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes:
+        """Wrap a key with the registered master key, without AAD, as Java and iceberg-rust do."""
+        return self._cipher(wrapping_key_id).encrypt(key)
+
+    def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes:
+        """Unwrap a key wrapped by `wrap_key`."""
+        return self._cipher(wrapping_key_id).decrypt(wrapped_key)

tests/encryption/test_kms.py

@@ -16,9 +16,10 @@
 # under the License.
 
 import pytest
+from memory_kms import MemoryKeyManagementClient
 
 from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey
-from pyiceberg.encryption.kms import GeneratedKey, KeyManagementClient, MemoryKeyManagementClient
+from pyiceberg.encryption.kms import GeneratedKey, KeyManagementClient
 
 MASTER_KEY_ID = "master-key"
 MASTER_KEY = SecureKey(b"0123456789012345")
@@ -60,7 +61,7 @@ def test_wrap_unwrap_round_trip(kms: MemoryKeyManagementClient) -> None:
 
 @pytest.mark.parametrize("key_size", list(AesKeySize))
 def test_wrap_unwrap_round_trip_for_each_master_key_size(key_size: AesKeySize) -> None:
-    kms = MemoryKeyManagementClient(key_size)
+    kms = MemoryKeyManagementClient(master_key_size=key_size)
     master_key = kms.add_master_key(MASTER_KEY_ID)
 
     assert master_key.key_size == key_size