From d23858809faf2c26c39c7b9f989bc776c14523d2 Mon Sep 17 00:00:00 2001 From: Xander Date: Fri, 11 Sep 2026 20:56:09 +0100 Subject: [PATCH 1/2] feat(encryption): [2/N] Add standard key metadata --- pyiceberg/encryption/__init__.py | 16 +++++ pyiceberg/encryption/key_metadata.py | 80 +++++++++++++++++++++++ tests/encryption/test_key_metadata.py | 93 +++++++++++++++++++++++++++ 3 files changed, 189 insertions(+) create mode 100644 pyiceberg/encryption/__init__.py create mode 100644 pyiceberg/encryption/key_metadata.py create mode 100644 tests/encryption/test_key_metadata.py diff --git a/pyiceberg/encryption/__init__.py b/pyiceberg/encryption/__init__.py new file mode 100644 index 0000000000..13a83393a9 --- /dev/null +++ b/pyiceberg/encryption/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/pyiceberg/encryption/key_metadata.py b/pyiceberg/encryption/key_metadata.py new file mode 100644 index 0000000000..d2173f7fa6 --- /dev/null +++ b/pyiceberg/encryption/key_metadata.py @@ -0,0 +1,80 @@ +# 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. +"""Key metadata for encrypted manifests, manifest lists and data files.""" + +from __future__ import annotations + +import io +from dataclasses import dataclass + +from pyiceberg.avro.decoder import new_decoder +from pyiceberg.avro.encoder import BinaryEncoder +from pyiceberg.avro.resolver import construct_reader, construct_writer +from pyiceberg.schema import Schema +from pyiceberg.typedef import Record +from pyiceberg.types import BinaryType, LongType, NestedField + +KEY_METADATA_V1 = 1 + +AES_KEY_LENGTHS = (16, 24, 32) + +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), + NestedField(field_id=2, name="file_length", field_type=LongType(), required=False), +) + + +@dataclass(frozen=True) +class StandardKeyMetadata: + """The key and AAD prefix needed to decrypt a single file. + + Wire format is a version byte followed by an Avro datum of `KEY_METADATA_SCHEMA_V1`, + byte-compatible with Java's `StandardKeyMetadata`. + """ + + encryption_key: bytes + aad_prefix: bytes | None = None + file_length: int | None = None + + def __post_init__(self) -> None: + """Reject invalid key lengths here rather than only on decode, so an invalid instance cannot exist.""" + if len(self.encryption_key) not in AES_KEY_LENGTHS: + raise ValueError( + f"Invalid encryption key in key metadata: expected one of {AES_KEY_LENGTHS} bytes, got {len(self.encryption_key)}" + ) + + @classmethod + def from_bytes(cls, data: bytes) -> StandardKeyMetadata: + """Decode key metadata from its wire format.""" + if not data: + raise ValueError("Empty key metadata") + + if (version := data[0]) != KEY_METADATA_V1: + raise ValueError(f"Unsupported key metadata version: {version}") + + record = construct_reader(KEY_METADATA_SCHEMA_V1).read(new_decoder(data[1:])) + return cls(encryption_key=record[0], aad_prefix=record[1], file_length=record[2]) + + def to_bytes(self) -> bytes: + """Encode key metadata to its wire format.""" + output = io.BytesIO() + encoder = BinaryEncoder(output) + encoder.write(bytes([KEY_METADATA_V1])) + record = Record(self.encryption_key, self.aad_prefix, self.file_length) + construct_writer(KEY_METADATA_SCHEMA_V1).write(encoder, record) + return output.getvalue() diff --git a/tests/encryption/test_key_metadata.py b/tests/encryption/test_key_metadata.py new file mode 100644 index 0000000000..786bb6c54e --- /dev/null +++ b/tests/encryption/test_key_metadata.py @@ -0,0 +1,93 @@ +# 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 io + +import pytest + +from pyiceberg.avro.encoder import BinaryEncoder +from pyiceberg.encryption.key_metadata import KEY_METADATA_V1, StandardKeyMetadata + +AES128_KEY = b"0123456789012345" + +# Version byte, then the Avro-encoded encryption_key: zigzag length 16 followed by the key +ENCODED_PREFIX = b"\x01\x20" + AES128_KEY + + +def encode_key_metadata(encryption_key: bytes) -> bytes: + """Encode key metadata directly, bypassing StandardKeyMetadata's validation.""" + output = io.BytesIO() + encoder = BinaryEncoder(output) + encoder.write(bytes([KEY_METADATA_V1])) + encoder.write_bytes(encryption_key) + encoder.write_int(0) + encoder.write_int(0) + return output.getvalue() + + +@pytest.mark.parametrize( + "key_metadata, encoded", + [ + (StandardKeyMetadata(encryption_key=AES128_KEY), ENCODED_PREFIX + b"\x00\x00"), + (StandardKeyMetadata(encryption_key=AES128_KEY, aad_prefix=b"ad"), ENCODED_PREFIX + b"\x02\x04ad\x00"), + ( + StandardKeyMetadata(encryption_key=AES128_KEY, aad_prefix=b"ad", file_length=1024), + ENCODED_PREFIX + b"\x02\x04ad\x02\x80\x10", + ), + (StandardKeyMetadata(encryption_key=AES128_KEY, aad_prefix=b""), ENCODED_PREFIX + b"\x02\x00\x00"), + ], +) +def test_key_metadata_serialization(key_metadata: StandardKeyMetadata, encoded: bytes) -> None: + assert key_metadata.to_bytes() == encoded + assert StandardKeyMetadata.from_bytes(encoded) == key_metadata + + +def test_key_metadata_defaults() -> None: + key_metadata = StandardKeyMetadata(encryption_key=AES128_KEY) + + assert key_metadata.aad_prefix is None + assert key_metadata.file_length is None + + +def test_key_metadata_empty_buffer() -> None: + with pytest.raises(ValueError, match="Empty key metadata"): + StandardKeyMetadata.from_bytes(b"") + + +@pytest.mark.parametrize("data", [b"\x02", b"\x02" + ENCODED_PREFIX[1:] + b"\x00\x00"]) +def test_key_metadata_unsupported_version(data: bytes) -> None: + with pytest.raises(ValueError, match="Unsupported key metadata version: 2"): + StandardKeyMetadata.from_bytes(data) + + +@pytest.mark.parametrize("key_length", [16, 24, 32]) +def test_key_metadata_accepts_aes_key_lengths(key_length: int) -> None: + key_metadata = StandardKeyMetadata(encryption_key=bytes(key_length)) + + assert StandardKeyMetadata.from_bytes(key_metadata.to_bytes()) == key_metadata + + +@pytest.mark.parametrize("key_length", [0, 4, 15, 20, 33]) +def test_key_metadata_rejects_invalid_key_length(key_length: int) -> None: + with pytest.raises(ValueError, match="Invalid encryption key in key metadata"): + StandardKeyMetadata(encryption_key=bytes(key_length)) + + +@pytest.mark.parametrize("key_length", [0, 4, 15, 20, 33]) +def test_key_metadata_decode_rejects_invalid_key_length(key_length: int) -> None: + with pytest.raises(ValueError, match="Invalid encryption key in key metadata"): + StandardKeyMetadata.from_bytes(encode_key_metadata(bytes(key_length))) From 6f6ceaab685d91f3af462add58f071d47305c5bd Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Fri, 11 Sep 2026 14:01:52 -0700 Subject: [PATCH 2/2] fix(encryption): Redact key metadata Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyiceberg/encryption/key_metadata.py | 4 ++-- tests/encryption/test_key_metadata.py | 7 +++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/pyiceberg/encryption/key_metadata.py b/pyiceberg/encryption/key_metadata.py index d2173f7fa6..f0c3d814cc 100644 --- a/pyiceberg/encryption/key_metadata.py +++ b/pyiceberg/encryption/key_metadata.py @@ -19,7 +19,7 @@ from __future__ import annotations import io -from dataclasses import dataclass +from dataclasses import dataclass, field from pyiceberg.avro.decoder import new_decoder from pyiceberg.avro.encoder import BinaryEncoder @@ -47,7 +47,7 @@ class StandardKeyMetadata: byte-compatible with Java's `StandardKeyMetadata`. """ - encryption_key: bytes + encryption_key: bytes = field(repr=False) aad_prefix: bytes | None = None file_length: int | None = None diff --git a/tests/encryption/test_key_metadata.py b/tests/encryption/test_key_metadata.py index 786bb6c54e..490f793c74 100644 --- a/tests/encryption/test_key_metadata.py +++ b/tests/encryption/test_key_metadata.py @@ -63,6 +63,13 @@ def test_key_metadata_defaults() -> None: assert key_metadata.file_length is None +def test_key_metadata_repr_redacts_encryption_key() -> None: + key_metadata = StandardKeyMetadata(encryption_key=AES128_KEY) + + assert "encryption_key" not in repr(key_metadata) + assert repr(AES128_KEY) not in repr(key_metadata) + + def test_key_metadata_empty_buffer() -> None: with pytest.raises(ValueError, match="Empty key metadata"): StandardKeyMetadata.from_bytes(b"")