feat(encryption): [8/N] encryption manager

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."""

pyiceberg/encryption/manager.py

@@ -0,0 +1,190 @@
+# 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."""
+
+    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")

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)