kevinjqliu commented on code in PR #4149: URL: https://github.com/apache/iceberg-python/pull/4149#discussion_r4371756496
########## 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)) Review Comment: nit: i think this and the other pytest.mark.parametrize can be simple for loops inside the function 😄 i added this rule in the AGENTS.md (see #4150) ########## 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): Review Comment: nit: `format_version < 3` might be better since we have v4 coming up soon. we can make this a follow up fix and refactor other places in the repo (if any) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
