This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 19abbcf1b feat(encryption): [8/N] encryption manager (#4149)
19abbcf1b is described below
commit 19abbcf1b01d60d9bd6b70a265b07e45811514b3
Author: Xander <[email protected]>
AuthorDate: Sun Oct 25 17:57:08 2026 +0000
feat(encryption): [8/N] encryption manager (#4149)
* feat(encryption): [8/N] encryption manager
* remove comment
* raise rather than None
* kek guard
* lazy import
---
pyiceberg/encryption/key_metadata.py | 9 +
pyiceberg/encryption/manager.py | 194 +++++++++++++++++++
tests/encryption/test_key_metadata.py | 11 ++
tests/encryption/test_manager.py | 337 ++++++++++++++++++++++++++++++++++
4 files changed, 551 insertions(+)
diff --git a/pyiceberg/encryption/key_metadata.py
b/pyiceberg/encryption/key_metadata.py
index f0c3d814c..e6546be26 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 @@ KEY_METADATA_V1 = 1
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 @@ class StandardKeyMetadata:
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 000000000..73c57dfd5
--- /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.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."""
+
+ 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.
+ """
+ from pyiceberg.table import TableProperties
+
+ if (table_key_id :=
metadata.properties.get(TableProperties.ENCRYPTION_KEY_ID)) is None:
+ return None
+
+ if not isinstance(metadata, TableMetadataV3):
+ raise ValueError(f"Invalid properties for
v{metadata.format_version}: {TableProperties.ENCRYPTION_KEY_ID}")
+
+ 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 key.encrypted_by_id == self._table_key_id:
+ raise ValueError(f"{key_id} is a key encryption key, not manifest
list key metadata")
+
+ 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 490f793c7..c01dea035 100644
--- a/tests/encryption/test_key_metadata.py
+++ b/tests/encryption/test_key_metadata.py
@@ -20,6 +20,7 @@ import io
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
[email protected]("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 000000000..0c282cdd9
--- /dev/null
+++ b/tests/encryption/test_manager.py
@@ -0,0 +1,337 @@
+# 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))
+
+
[email protected]
+def kms() -> CountingKeyManagementClient:
+ client = CountingKeyManagementClient()
+ client.add_master_key(MASTER_KEY_ID)
+ return client
+
+
[email protected]
+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
+
+
[email protected]("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 test_unwrap_kek_id(manager: EncryptionManager) -> None:
+ manager.encrypt_manifest_list_key_metadata(KEY_METADATA)
+ kek = manager.encryption_keys[0]
+
+ with pytest.raises(ValueError, match=f"{kek.key_id} is a key encryption
key, not manifest list key metadata"):
+ manager.decrypt_manifest_list_key_metadata(kek.key_id)
+
+
+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)
+
+
[email protected]("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}})
+
+ with pytest.raises(ValueError, match="Invalid properties for v2:
encryption.key-id"):
+ EncryptionManager.from_table_metadata(kms, metadata)
+
+
+def test_from_table_metadata_before_v3_without_key_id(
+ kms: CountingKeyManagementClient, example_table_metadata_v2: dict[str, Any]
+) -> None:
+ metadata = TableMetadataUtil.parse_obj({**example_table_metadata_v2,
"properties": {}})
+
+ 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)
