From 2cd35efee32387f772e7aa721a1df2e639548ce0 Mon Sep 17 00:00:00 2001 From: Xander Date: Sat, 24 Oct 2026 14:02:27 -0400 Subject: [PATCH 1/2] feat(encryption): [8/N] encryption manager --- pyiceberg/encryption/key_metadata.py | 9 + pyiceberg/encryption/manager.py | 194 ++++++++++++++++ tests/encryption/test_key_metadata.py | 11 + tests/encryption/test_manager.py | 320 ++++++++++++++++++++++++++ 4 files changed, 534 insertions(+) create mode 100644 pyiceberg/encryption/manager.py create mode 100644 tests/encryption/test_manager.py diff --git a/pyiceberg/encryption/key_metadata.py b/pyiceberg/encryption/key_metadata.py index f0c3d814cc..e6546be26f 100644 --- a/pyiceberg/encryption/key_metadata.py +++ b/pyiceberg/encryption/key_metadata.py @@ -19,11 +19,13 @@ from __future__ import annotations import io +import os from dataclasses import dataclass, field from pyiceberg.avro.decoder import new_decoder from pyiceberg.avro.encoder import BinaryEncoder from pyiceberg.avro.resolver import construct_reader, construct_writer +from pyiceberg.encryption.ciphers import AesKeySize, SecureKey from pyiceberg.schema import Schema from pyiceberg.typedef import Record from pyiceberg.types import BinaryType, LongType, NestedField @@ -32,6 +34,8 @@ AES_KEY_LENGTHS = (16, 24, 32) +AAD_PREFIX_LENGTH = 16 + KEY_METADATA_SCHEMA_V1 = Schema( NestedField(field_id=0, name="encryption_key", field_type=BinaryType(), required=True), NestedField(field_id=1, name="aad_prefix", field_type=BinaryType(), required=False), @@ -58,6 +62,11 @@ def __post_init__(self) -> None: f"Invalid encryption key in key metadata: expected one of {AES_KEY_LENGTHS} bytes, got {len(self.encryption_key)}" ) + @classmethod + def generate(cls, key_size: AesKeySize = AesKeySize.BITS_128) -> StandardKeyMetadata: + """Generate key metadata for a new file, with a random key of `key_size` and a random AAD prefix.""" + return cls(encryption_key=SecureKey.generate(key_size).key, aad_prefix=os.urandom(AAD_PREFIX_LENGTH)) + @classmethod def from_bytes(cls, data: bytes) -> StandardKeyMetadata: """Decode key metadata from its wire format.""" diff --git a/pyiceberg/encryption/manager.py b/pyiceberg/encryption/manager.py new file mode 100644 index 0000000000..d3f0ef939f --- /dev/null +++ b/pyiceberg/encryption/manager.py @@ -0,0 +1,194 @@ +# 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. +"""Envelope key management for table encryption. + +A master key held by the key management service wraps key encryption keys (KEKs), which are +stored in table metadata. A KEK wraps the key metadata of each snapshot's manifest list. +""" + +from __future__ import annotations + +import time +import uuid +from collections.abc import Iterable + +from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey +from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.kms import KeyManagementClient +from pyiceberg.table import TableProperties +from pyiceberg.table.metadata import EncryptedKey, TableMetadata, TableMetadataV3 +from pyiceberg.utils.properties import property_as_int + +KEK_CREATED_AT_PROPERTY = "KEY_TIMESTAMP" + +KEK_LIFESPAN_MS = 730 * 24 * 60 * 60 * 1000 + +KEK_CACHE_TTL_SECONDS = 3600 + + +class EncryptionManager: + """Generates file keys and wraps manifest list key metadata for an encrypted table. + + Mirrors iceberg-rust's `EncryptionManager`. Keys it writes can be read by Java's + `StandardEncryptionManager`, and the reverse. + """ + + def __init__( + self, + kms_client: KeyManagementClient, + table_key_id: str, + encryption_keys: Iterable[EncryptedKey] = (), + key_size: AesKeySize = AesKeySize.BITS_128, + ) -> None: + self._kms_client = kms_client + self._table_key_id = table_key_id + self._key_size = key_size + self._encryption_keys = {key.key_id: key for key in encryption_keys} + self._kek_cache: dict[str, tuple[SecureKey, float]] = {} + + @classmethod + def from_table_metadata(cls, kms_client: KeyManagementClient | None, metadata: TableMetadata) -> EncryptionManager | None: + """Create a manager for the table, or return None if the table is not encrypted. + + Args: + kms_client (KeyManagementClient | None): Wraps and unwraps KEKs with the table's master key. + metadata (TableMetadata): The table's metadata, holding `encryption.key-id` and its encryption keys. + """ + if not isinstance(metadata, TableMetadataV3): + return None + + if (table_key_id := metadata.properties.get(TableProperties.ENCRYPTION_KEY_ID)) is None: + return None + + if kms_client is None: + raise ValueError(f"Table has {TableProperties.ENCRYPTION_KEY_ID} set but no key management client was provided") + + key_length = property_as_int(metadata.properties, TableProperties.ENCRYPTION_DATA_KEY_LENGTH) + key_size = AesKeySize.from_key_length( + TableProperties.ENCRYPTION_DATA_KEY_LENGTH_DEFAULT if key_length is None else key_length + ) + + return cls(kms_client, table_key_id, metadata.encryption_keys, key_size) + + @property + def encryption_keys(self) -> list[EncryptedKey]: + """The table's encryption keys, including KEKs and manifest list keys this manager created.""" + return list(self._encryption_keys.values()) + + def generate_key_metadata(self) -> StandardKeyMetadata: + """Generate key metadata for a new file, with a fresh key and AAD prefix.""" + return StandardKeyMetadata.generate(self._key_size) + + def encrypt_manifest_list_key_metadata(self, key_metadata: StandardKeyMetadata) -> str: + """Wrap a manifest list's key metadata with the active KEK, creating a KEK if none is active. + + Returns the id of the new encryption key, to record as the snapshot's `key-id`. + + Args: + key_metadata (StandardKeyMetadata): The manifest list's key metadata, including its `file_length`. + """ + kek = self._find_active_kek() or self._create_kek() + wrapped = AesGcmCipher(self._unwrap_kek(kek)).encrypt(key_metadata.to_bytes(), _kek_timestamp_aad(kek)) + + key = EncryptedKey(key_id=str(uuid.uuid4()), encrypted_key_metadata=wrapped, encrypted_by_id=kek.key_id) + self._encryption_keys[key.key_id] = key + return key.key_id + + def decrypt_manifest_list_key_metadata(self, key_id: str) -> StandardKeyMetadata: + """Unwrap the manifest list key metadata stored under `key_id`. + + Args: + key_id (str): The snapshot's `key-id`. + """ + if (key := self._encryption_keys.get(key_id)) is None: + raise ValueError(f"Encryption key not found: {key_id}") + + if key.encrypted_by_id is None: + raise ValueError(f"Encryption key {key_id} has no encrypted-by-id") + + if (kek := self._encryption_keys.get(key.encrypted_by_id)) is None: + raise ValueError(f"KEK not found in encryption keys: {key.encrypted_by_id}") + + aad = _kek_timestamp_aad(kek) + cipher = AesGcmCipher(self._unwrap_kek(kek)) + try: + plaintext = cipher.decrypt(key.encrypted_key_metadata, aad) + except ValueError as e: + raise ValueError(f"Failed to unwrap key metadata with KEK {kek.key_id}") from e + + return StandardKeyMetadata.from_bytes(plaintext) + + def _create_kek(self) -> EncryptedKey: + if self._kms_client.supports_key_generation(): + generated = self._kms_client.generate_key(self._table_key_id) + plaintext, wrapped = SecureKey(generated.key), generated.wrapped_key + else: + plaintext = SecureKey.generate(self._key_size) + wrapped = self._kms_client.wrap_key(plaintext.key, self._table_key_id) + + kek = EncryptedKey( + key_id=str(uuid.uuid4()), + encrypted_key_metadata=wrapped, + encrypted_by_id=self._table_key_id, + properties={KEK_CREATED_AT_PROPERTY: str(_now_ms())}, + ) + self._kek_cache[kek.key_id] = (plaintext, time.monotonic()) + self._encryption_keys[kek.key_id] = kek + return kek + + def _find_active_kek(self) -> EncryptedKey | None: + keks = [ + key + for key in self._encryption_keys.values() + if key.encrypted_by_id == self._table_key_id and not _is_kek_expired(key) + ] + return max(keks, key=lambda kek: _kek_created_at(kek) or 0, default=None) + + def _unwrap_kek(self, kek: EncryptedKey) -> SecureKey: + if (cached := self._kek_cache.get(kek.key_id)) is not None and time.monotonic() - cached[1] < KEK_CACHE_TTL_SECONDS: + return cached[0] + + if kek.encrypted_by_id is None: + raise ValueError(f"KEK {kek.key_id} has no encrypted-by-id") + + plaintext = SecureKey(self._kms_client.unwrap_key(kek.encrypted_key_metadata, kek.encrypted_by_id)) + self._kek_cache[kek.key_id] = (plaintext, time.monotonic()) + return plaintext + + +def _now_ms() -> int: + return time.time_ns() // 1_000_000 + + +def _kek_created_at(kek: EncryptedKey) -> int | None: + try: + return int(kek.properties[KEK_CREATED_AT_PROPERTY]) + except (KeyError, ValueError): + return None + + +def _is_kek_expired(kek: EncryptedKey) -> bool: + created_at = _kek_created_at(kek) + return created_at is None or _now_ms() - created_at >= KEK_LIFESPAN_MS + + +def _kek_timestamp_aad(kek: EncryptedKey) -> bytes: + """Return the KEK's creation timestamp, which authenticates everything it wraps so the timestamp cannot be altered.""" + if (timestamp := kek.properties.get(KEK_CREATED_AT_PROPERTY)) is None: + raise ValueError(f"KEK {kek.key_id} is missing required {KEK_CREATED_AT_PROPERTY} property") + + return timestamp.encode("utf-8") diff --git a/tests/encryption/test_key_metadata.py b/tests/encryption/test_key_metadata.py index 490f793c74..c01dea0359 100644 --- a/tests/encryption/test_key_metadata.py +++ b/tests/encryption/test_key_metadata.py @@ -20,6 +20,7 @@ import pytest from pyiceberg.avro.encoder import BinaryEncoder +from pyiceberg.encryption.ciphers import AesKeySize from pyiceberg.encryption.key_metadata import KEY_METADATA_V1, StandardKeyMetadata AES128_KEY = b"0123456789012345" @@ -63,6 +64,16 @@ def test_key_metadata_defaults() -> None: assert key_metadata.file_length is None +@pytest.mark.parametrize("key_size", list(AesKeySize)) +def test_key_metadata_generate(key_size: AesKeySize) -> None: + key_metadata = StandardKeyMetadata.generate(key_size) + + assert len(key_metadata.encryption_key) == key_size.key_length + assert key_metadata.aad_prefix is not None and len(key_metadata.aad_prefix) == 16 + assert key_metadata.file_length is None + assert StandardKeyMetadata.generate(key_size) != key_metadata + + def test_key_metadata_repr_redacts_encryption_key() -> None: key_metadata = StandardKeyMetadata(encryption_key=AES128_KEY) diff --git a/tests/encryption/test_manager.py b/tests/encryption/test_manager.py new file mode 100644 index 0000000000..f46f6737b5 --- /dev/null +++ b/tests/encryption/test_manager.py @@ -0,0 +1,320 @@ +# 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. +import time +from typing import Any + +import pytest +from memory_kms import MemoryKeyManagementClient + +from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey +from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.kms import GeneratedKey +from pyiceberg.encryption.manager import KEK_CREATED_AT_PROPERTY, EncryptionManager +from pyiceberg.table.metadata import EncryptedKey, TableMetadataUtil + +MASTER_KEY_ID = "master-1" +KEY_METADATA = StandardKeyMetadata(encryption_key=b"0123456789abcdef", aad_prefix=b"test-aad-prefix!", file_length=1024) +THREE_YEARS_MS = 3 * 365 * 24 * 60 * 60 * 1000 + + +class CountingKeyManagementClient(MemoryKeyManagementClient): + unwrap_calls = 0 + + def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes: + self.unwrap_calls += 1 + return super().unwrap_key(wrapped_key, wrapping_key_id) + + +class KeyGeneratingKeyManagementClient(MemoryKeyManagementClient): + def supports_key_generation(self) -> bool: + return True + + def generate_key(self, wrapping_key_id: str) -> GeneratedKey: + key = SecureKey.generate().key + return GeneratedKey(key=key, wrapped_key=self.wrap_key(key, wrapping_key_id)) + + +@pytest.fixture +def kms() -> CountingKeyManagementClient: + client = CountingKeyManagementClient() + client.add_master_key(MASTER_KEY_ID) + return client + + +@pytest.fixture +def manager(kms: CountingKeyManagementClient) -> EncryptionManager: + return EncryptionManager(kms, MASTER_KEY_ID) + + +def kek_created_at(kms: MemoryKeyManagementClient, key_id: str, created_at_ms: int) -> EncryptedKey: + return EncryptedKey( + key_id=key_id, + encrypted_key_metadata=kms.wrap_key(SecureKey.generate().key, MASTER_KEY_ID), + encrypted_by_id=MASTER_KEY_ID, + properties={KEK_CREATED_AT_PROPERTY: str(created_at_ms)}, + ) + + +def now_ms() -> int: + return time.time_ns() // 1_000_000 + + +def test_first_wrap_creates_a_kek(manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + kek, manifest_list_key = manager.encryption_keys + assert kek.encrypted_by_id == MASTER_KEY_ID + assert kek.encrypted_key_metadata + assert abs(int(kek.properties[KEK_CREATED_AT_PROPERTY]) - now_ms()) < 60_000 + assert manifest_list_key.key_id == key_id + assert manifest_list_key.encrypted_by_id == kek.key_id + assert manifest_list_key.properties == {} + + +def test_wrap_unwrap_round_trip(manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + assert manager.decrypt_manifest_list_key_metadata(key_id) == KEY_METADATA + + +def test_unwrap_with_keys_from_another_manager(kms: CountingKeyManagementClient, manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + reader = EncryptionManager(kms, MASTER_KEY_ID, manager.encryption_keys) + + assert reader.decrypt_manifest_list_key_metadata(key_id) == KEY_METADATA + + +def test_unwrap_key_metadata_wrapped_as_java_does(kms: CountingKeyManagementClient) -> None: + """Java's `EncryptionUtil.encryptManifestListKeyMetadata` wraps with AES-GCM, using the KEK timestamp as AAD.""" + kek_key = SecureKey.generate() + kek = EncryptedKey( + key_id="kek", + encrypted_key_metadata=kms.wrap_key(kek_key.key, MASTER_KEY_ID), + encrypted_by_id=MASTER_KEY_ID, + properties={KEK_CREATED_AT_PROPERTY: "1700000000000"}, + ) + manifest_list_key = EncryptedKey( + key_id="manifest-list-key", + encrypted_key_metadata=AesGcmCipher(kek_key).encrypt(KEY_METADATA.to_bytes(), b"1700000000000"), + encrypted_by_id="kek", + ) + + manager = EncryptionManager(kms, MASTER_KEY_ID, [kek, manifest_list_key]) + + assert manager.decrypt_manifest_list_key_metadata("manifest-list-key") == KEY_METADATA + + +def test_kek_is_reused_while_active(manager: EncryptionManager) -> None: + first = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + second = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + kek, first_key, second_key = manager.encryption_keys + assert (first_key.key_id, second_key.key_id) == (first, second) + assert first_key.encrypted_by_id == second_key.encrypted_by_id == kek.key_id + + +def test_kek_is_rotated_once_expired(kms: CountingKeyManagementClient) -> None: + expired = kek_created_at(kms, "expired-kek", now_ms() - THREE_YEARS_MS) + manager = EncryptionManager(kms, MASTER_KEY_ID, [expired]) + + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + old_kek, new_kek, manifest_list_key = manager.encryption_keys + assert old_kek == expired + assert new_kek.key_id != expired.key_id + assert manifest_list_key.encrypted_by_id == new_kek.key_id + assert manager.decrypt_manifest_list_key_metadata(key_id) == KEY_METADATA + + +@pytest.mark.parametrize("properties", [{}, {KEK_CREATED_AT_PROPERTY: "not-a-timestamp"}]) +def test_kek_without_a_valid_timestamp_is_not_reused(kms: CountingKeyManagementClient, properties: dict[str, str]) -> None: + kek = EncryptedKey( + key_id="kek", + encrypted_key_metadata=kms.wrap_key(SecureKey.generate().key, MASTER_KEY_ID), + encrypted_by_id=MASTER_KEY_ID, + properties=properties, + ) + manager = EncryptionManager(kms, MASTER_KEY_ID, [kek]) + + manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + _, new_kek, manifest_list_key = manager.encryption_keys + assert new_kek.encrypted_by_id == MASTER_KEY_ID + assert manifest_list_key.encrypted_by_id == new_kek.key_id + + +def test_newest_active_kek_is_used(kms: CountingKeyManagementClient) -> None: + newer = kek_created_at(kms, "newer", now_ms() - 1_000) + older = kek_created_at(kms, "older", now_ms() - 2_000) + manager = EncryptionManager(kms, MASTER_KEY_ID, [newer, older]) + + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + assert len(manager.encryption_keys) == 3 + assert manager.encryption_keys[-1].key_id == key_id + assert manager.encryption_keys[-1].encrypted_by_id == "newer" + + +def test_kek_wrapped_by_another_master_key_is_not_used(kms: CountingKeyManagementClient) -> None: + kms.add_master_key("master-2") + other = EncryptedKey( + key_id="other", + encrypted_key_metadata=kms.wrap_key(SecureKey.generate().key, "master-2"), + encrypted_by_id="master-2", + properties={KEK_CREATED_AT_PROPERTY: str(now_ms())}, + ) + manager = EncryptionManager(kms, MASTER_KEY_ID, [other]) + + manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + _, new_kek, manifest_list_key = manager.encryption_keys + assert new_kek.encrypted_by_id == MASTER_KEY_ID + assert manifest_list_key.encrypted_by_id == new_kek.key_id + + +def test_kek_is_generated_by_the_kms_when_supported() -> None: + kms = KeyGeneratingKeyManagementClient() + kms.add_master_key(MASTER_KEY_ID) + manager = EncryptionManager(kms, MASTER_KEY_ID) + + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + + assert ( + EncryptionManager(kms, MASTER_KEY_ID, manager.encryption_keys).decrypt_manifest_list_key_metadata(key_id) == KEY_METADATA + ) + + +def test_unwrapped_kek_is_cached(kms: CountingKeyManagementClient, manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + manager.decrypt_manifest_list_key_metadata(key_id) + assert kms.unwrap_calls == 0 + + reader = EncryptionManager(kms, MASTER_KEY_ID, manager.encryption_keys) + reader.decrypt_manifest_list_key_metadata(key_id) + reader.decrypt_manifest_list_key_metadata(key_id) + assert kms.unwrap_calls == 1 + + +def test_unwrap_unknown_key_id(manager: EncryptionManager) -> None: + with pytest.raises(ValueError, match="Encryption key not found: missing"): + manager.decrypt_manifest_list_key_metadata("missing") + + +def test_unwrap_key_without_encrypted_by_id(kms: CountingKeyManagementClient) -> None: + manager = EncryptionManager(kms, MASTER_KEY_ID, [EncryptedKey(key_id="key", encrypted_key_metadata=b"wrapped")]) + + with pytest.raises(ValueError, match="Encryption key key has no encrypted-by-id"): + manager.decrypt_manifest_list_key_metadata("key") + + +def test_unwrap_key_with_missing_kek(kms: CountingKeyManagementClient) -> None: + key = EncryptedKey(key_id="key", encrypted_key_metadata=b"wrapped", encrypted_by_id="missing-kek") + manager = EncryptionManager(kms, MASTER_KEY_ID, [key]) + + with pytest.raises(ValueError, match="KEK not found in encryption keys: missing-kek"): + manager.decrypt_manifest_list_key_metadata("key") + + +def replace_kek_properties(manager: EncryptionManager, key_id: str, properties: dict[str, str]) -> list[EncryptedKey]: + kek_id = next(key.encrypted_by_id for key in manager.encryption_keys if key.key_id == key_id) + return [key.model_copy(update={"properties": properties}) if key.key_id == kek_id else key for key in manager.encryption_keys] + + +def test_unwrap_fails_when_kek_timestamp_is_missing(kms: CountingKeyManagementClient, manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + reader = EncryptionManager(kms, MASTER_KEY_ID, replace_kek_properties(manager, key_id, {})) + + with pytest.raises(ValueError, match=f"is missing required {KEK_CREATED_AT_PROPERTY} property"): + reader.decrypt_manifest_list_key_metadata(key_id) + + +def test_unwrap_fails_when_kek_timestamp_is_tampered(kms: CountingKeyManagementClient, manager: EncryptionManager) -> None: + key_id = manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + reader = EncryptionManager(kms, MASTER_KEY_ID, replace_kek_properties(manager, key_id, {KEK_CREATED_AT_PROPERTY: "9999999"})) + + with pytest.raises(ValueError, match="Failed to unwrap key metadata with KEK"): + reader.decrypt_manifest_list_key_metadata(key_id) + + +@pytest.mark.parametrize("key_size", list(AesKeySize)) +def test_generate_key_metadata(kms: CountingKeyManagementClient, key_size: AesKeySize) -> None: + key_metadata = EncryptionManager(kms, MASTER_KEY_ID, key_size=key_size).generate_key_metadata() + + assert len(key_metadata.encryption_key) == key_size.key_length + + +def test_from_table_metadata(kms: CountingKeyManagementClient, example_table_metadata_v3: dict[str, Any]) -> None: + kek = kek_created_at(kms, "kek", now_ms()) + metadata = TableMetadataUtil.parse_obj( + { + **example_table_metadata_v3, + "properties": {"encryption.key-id": MASTER_KEY_ID, "encryption.data-key-length": "32"}, + "encryption-keys": [kek.model_dump(by_alias=True, mode="json")], + } + ) + + manager = EncryptionManager.from_table_metadata(kms, metadata) + + assert manager is not None + assert manager.encryption_keys == [kek] + assert len(manager.generate_key_metadata().encryption_key) == 32 + manager.encrypt_manifest_list_key_metadata(KEY_METADATA) + assert manager.encryption_keys[-1].encrypted_by_id == "kek" + + +def test_from_table_metadata_defaults_to_128_bit_keys( + kms: CountingKeyManagementClient, example_table_metadata_v3: dict[str, Any] +) -> None: + metadata = TableMetadataUtil.parse_obj({**example_table_metadata_v3, "properties": {"encryption.key-id": MASTER_KEY_ID}}) + + manager = EncryptionManager.from_table_metadata(kms, metadata) + + assert manager is not None + assert len(manager.generate_key_metadata().encryption_key) == 16 + + +def test_from_table_metadata_without_key_id(kms: CountingKeyManagementClient, example_table_metadata_v3: dict[str, Any]) -> None: + metadata = TableMetadataUtil.parse_obj({**example_table_metadata_v3, "properties": {}}) + + assert EncryptionManager.from_table_metadata(kms, metadata) is None + assert EncryptionManager.from_table_metadata(None, metadata) is None + + +def test_from_table_metadata_before_v3(kms: CountingKeyManagementClient, example_table_metadata_v2: dict[str, Any]) -> None: + metadata = TableMetadataUtil.parse_obj({**example_table_metadata_v2, "properties": {"encryption.key-id": MASTER_KEY_ID}}) + + assert EncryptionManager.from_table_metadata(kms, metadata) is None + + +def test_from_table_metadata_without_kms_client(example_table_metadata_v3: dict[str, Any]) -> None: + metadata = TableMetadataUtil.parse_obj({**example_table_metadata_v3, "properties": {"encryption.key-id": MASTER_KEY_ID}}) + + with pytest.raises(ValueError, match="Table has encryption.key-id set but no key management client was provided"): + EncryptionManager.from_table_metadata(None, metadata) + + +def test_from_table_metadata_with_invalid_data_key_length( + kms: CountingKeyManagementClient, example_table_metadata_v3: dict[str, Any] +) -> None: + metadata = TableMetadataUtil.parse_obj( + {**example_table_metadata_v3, "properties": {"encryption.key-id": MASTER_KEY_ID, "encryption.data-key-length": "20"}} + ) + + with pytest.raises(ValueError, match="Unsupported key length: 20"): + EncryptionManager.from_table_metadata(kms, metadata) From 57994ae6c4f419bf05397dcdc273f9f352a021e4 Mon Sep 17 00:00:00 2001 From: Xander Date: Sat, 24 Oct 2026 14:13:15 -0400 Subject: [PATCH 2/2] remove comment --- pyiceberg/encryption/manager.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pyiceberg/encryption/manager.py b/pyiceberg/encryption/manager.py index d3f0ef939f..35877d46f3 100644 --- a/pyiceberg/encryption/manager.py +++ b/pyiceberg/encryption/manager.py @@ -41,11 +41,7 @@ class EncryptionManager: - """Generates file keys and wraps manifest list key metadata for an encrypted table. - - Mirrors iceberg-rust's `EncryptionManager`. Keys it writes can be read by Java's - `StandardEncryptionManager`, and the reverse. - """ + """Generates file keys and wraps manifest list key metadata for an encrypted table.""" def __init__( self,