drop REST mentions and trim verbose comments/docstrings

Author: smaheshwar-pltr

dev/provision.py

@@ -396,9 +396,7 @@
     spark.sql(f"ALTER TABLE {catalog_name}.default.test_empty_scan_ordered_str WRITE ORDERED BY id")
     spark.sql(f"INSERT INTO {catalog_name}.default.test_empty_scan_ordered_str VALUES 'a', 'c'")
 
-# Encrypted Iceberg table written via Spark, read back via PyIceberg in tests/integration/test_encryption.py.
-# Only the Hive catalog is configured with a Java-side KMS (encryption.kms-impl=UnitestKMS); the REST catalog
-# image does not ship UnitestKMS so we limit this fixture to Hive.
+# Encrypted Hive-cataloged table; read back via PyIceberg in tests/integration/test_encryption.py.
 spark.sql("""
     CREATE OR REPLACE TABLE hive.default.test_encrypted (id bigint, data string, value float)
     USING iceberg

pyiceberg/encryption/ciphers.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""AES-GCM encryption/decryption primitives and AGS1 stream decryption."""
+"""AES-GCM primitives and Iceberg AGS1 stream decryption."""
 
 from __future__ import annotations
 
@@ -28,58 +28,37 @@
 
 
 def aes_gcm_encrypt(key: bytes, plaintext: bytes, aad: bytes | None = None) -> bytes:
-    """Encrypt using AES-GCM. Returns nonce || ciphertext || tag."""
     nonce = os.urandom(NONCE_LENGTH)
-    aesgcm = AESGCM(key)
-    ciphertext_with_tag = aesgcm.encrypt(nonce, plaintext, aad)
-    return nonce + ciphertext_with_tag
+    return nonce + AESGCM(key).encrypt(nonce, plaintext, aad)
 
 
 def aes_gcm_decrypt(key: bytes, ciphertext: bytes, aad: bytes | None = None) -> bytes:
-    """Decrypt AES-GCM data in format: nonce || ciphertext || tag."""
     if len(ciphertext) < NONCE_LENGTH + GCM_TAG_LENGTH:
         raise ValueError(f"Ciphertext too short: {len(ciphertext)} bytes")
-    nonce = ciphertext[:NONCE_LENGTH]
-    encrypted_data = ciphertext[NONCE_LENGTH:]
-    aesgcm = AESGCM(key)
-    return aesgcm.decrypt(nonce, encrypted_data, aad)
+    return AESGCM(key).decrypt(ciphertext[:NONCE_LENGTH], ciphertext[NONCE_LENGTH:], aad)
 
 
-# AGS1 stream constants
 GCM_STREAM_MAGIC = b"AGS1"
-GCM_STREAM_HEADER_LENGTH = 8  # 4 magic + 4 block size
+GCM_STREAM_HEADER_LENGTH = 8  # 4 magic + 4 little-endian block size
 
 
 def stream_block_aad(aad_prefix: bytes, block_index: int) -> bytes:
-    """Construct per-block AAD for AGS1 stream encryption.
-
-    Format: aad_prefix || block_index (4 bytes, little-endian).
-    """
-    index_bytes = struct.pack("<I", block_index)
-    if not aad_prefix:
-        return index_bytes
-    return aad_prefix + index_bytes
+    return aad_prefix + struct.pack("<I", block_index)
 
 
 def decrypt_ags1_stream(key: bytes, encrypted_data: bytes, aad_prefix: bytes) -> bytes:
-    """Decrypt an entire AGS1 stream and return the plaintext.
-
-    AGS1 format:
-      - Header: "AGS1" (4 bytes) + plain_block_size (4 bytes LE)
-      - Blocks: each block is nonce(12) + ciphertext(up to 1MB) + tag(16)
-      - Each block's AAD = aad_prefix + block_index (4 bytes LE)
+    """Decrypt an Iceberg AGS1 stream.
 
+    Layout: "AGS1" (4) | plain_block_size LE (4) | one or more {nonce(12) | cipher | tag(16)} blocks.
+    Each block's AAD is `aad_prefix || block_index_le32`.
     """
     if len(encrypted_data) < GCM_STREAM_HEADER_LENGTH:
         raise ValueError(f"AGS1 stream too short: {len(encrypted_data)} bytes")
-
-    magic = encrypted_data[:4]
-    if magic != GCM_STREAM_MAGIC:
-        raise ValueError(f"Invalid AGS1 magic: {magic!r}, expected {GCM_STREAM_MAGIC!r}")
+    if encrypted_data[:4] != GCM_STREAM_MAGIC:
+        raise ValueError(f"Invalid AGS1 magic: {encrypted_data[:4]!r}")
 
     plain_block_size = struct.unpack_from("<I", encrypted_data, 4)[0]
     cipher_block_size = plain_block_size + NONCE_LENGTH + GCM_TAG_LENGTH
-
     stream_data = encrypted_data[GCM_STREAM_HEADER_LENGTH:]
     if not stream_data:
         return b""
@@ -88,28 +67,15 @@ def decrypt_ags1_stream(key: bytes, encrypted_data: bytes, aad_prefix: bytes) ->
     result = bytearray()
     offset = 0
     block_index = 0
-
     while offset < len(stream_data):
-        # Determine this block's cipher size
-        remaining = len(stream_data) - offset
-        if remaining >= cipher_block_size:
-            block_cipher_size = cipher_block_size
-        else:
-            block_cipher_size = remaining
-
+        block_cipher_size = min(cipher_block_size, len(stream_data) - offset)
         if block_cipher_size < NONCE_LENGTH + GCM_TAG_LENGTH:
-            raise ValueError(
-                f"Truncated AGS1 block at offset {offset}: {block_cipher_size} bytes (minimum {NONCE_LENGTH + GCM_TAG_LENGTH})"
-            )
-
-        block_data = stream_data[offset : offset + block_cipher_size]
-        nonce = block_data[:NONCE_LENGTH]
-        ciphertext_with_tag = block_data[NONCE_LENGTH:]
-
-        aad = stream_block_aad(aad_prefix, block_index)
-        plaintext = aesgcm.decrypt(nonce, ciphertext_with_tag, aad)
-        result.extend(plaintext)
+            raise ValueError(f"Truncated AGS1 block at offset {offset}: {block_cipher_size} bytes")
 
+        block = stream_data[offset : offset + block_cipher_size]
+        result.extend(
+            aesgcm.decrypt(block[:NONCE_LENGTH], block[NONCE_LENGTH:], stream_block_aad(aad_prefix, block_index))
+        )
         offset += block_cipher_size
         block_index += 1
 

pyiceberg/encryption/io.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""InputFile implementation backed by in-memory bytes."""
+"""In-memory InputFile/InputStream used to wrap decrypted Avro buffers for AvroFile."""
 
 from __future__ import annotations
 
@@ -25,8 +25,6 @@
 
 
 class BytesInputStream(InputStream):
-    """InputStream implementation backed by a bytes buffer."""
-
     def __init__(self, data: bytes) -> None:
         self._buffer = io.BytesIO(data)
 
@@ -45,7 +43,6 @@ def close(self) -> None:
         self._buffer.close()
 
     def __enter__(self) -> BytesInputStream:
-        """Enter the context manager."""
         return self
 
     def __exit__(
@@ -54,23 +51,15 @@ def __exit__(
         excinst: BaseException | None,
         exctb: TracebackType | None,
     ) -> None:
-        """Exit the context manager and close the stream."""
         self.close()
 
 
 class BytesInputFile(InputFile):
-    """InputFile implementation backed by in-memory bytes.
-
-    Used to wrap decrypted data so that it can be read by
-    AvroFile and other readers that expect an InputFile.
-    """
-
     def __init__(self, location: str, data: bytes) -> None:
         super().__init__(location)
         self._data = data
 
     def __len__(self) -> int:
-        """Return the length of the underlying data."""
         return len(self._data)
 
     def exists(self) -> bool:

pyiceberg/encryption/key_metadata.py

@@ -14,14 +14,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""StandardKeyMetadata Avro serialization.
+"""StandardKeyMetadata Avro codec.
 
-Wire format: ``0x01 version byte || Avro-encoded fields``
-
-Avro schema:
-  - encryption_key: bytes (required)
-  - aad_prefix: union[null, bytes] (optional)
-  - file_length: union[null, long] (optional)
+Wire: ``0x01 version`` || encryption_key (bytes) || aad_prefix (union[null,bytes])
+                       || file_length (union[null,long]).
 """
 
 from __future__ import annotations
@@ -32,7 +28,6 @@
 
 
 def _read_avro_long(data: bytes, offset: int) -> tuple[int, int]:
-    """Read a zigzag-encoded Avro long from data at offset. Returns (value, new_offset)."""
     result = 0
     shift = 0
     while True:
@@ -44,12 +39,10 @@ def _read_avro_long(data: bytes, offset: int) -> tuple[int, int]:
         if (b & 0x80) == 0:
             break
         shift += 7
-    # Zigzag decode
     return (result >> 1) ^ -(result & 1), offset
 
 
 def _read_avro_bytes(data: bytes, offset: int) -> tuple[bytes, int]:
-    """Read Avro bytes (length-prefixed). Returns (bytes_value, new_offset)."""
     length, offset = _read_avro_long(data, offset)
     if length < 0:
         raise ValueError(f"Negative Avro bytes length: {length}")
@@ -61,31 +54,20 @@ def _read_avro_bytes(data: bytes, offset: int) -> tuple[bytes, int]:
 
 @dataclass(frozen=True)
 class StandardKeyMetadata:
-    """Standard key metadata for Iceberg table encryption.
-
-    Contains the plaintext encryption key (DEK), AAD prefix, and optional file length.
-    """
-
     encryption_key: bytes
     aad_prefix: bytes = b""
     file_length: int | None = None
 
     @staticmethod
     def deserialize(data: bytes) -> StandardKeyMetadata:
-        """Deserialize from wire format: ``0x01 version || Avro-encoded fields``."""
         if not data:
             raise ValueError("Empty key metadata buffer")
-
-        version = data[0]
-        if version != V1:
-            raise ValueError(f"Unsupported key metadata version: {version}")
-
+        if data[0] != V1:
+            raise ValueError(f"Unsupported key metadata version: {data[0]}")
         offset = 1
 
-        # Read encryption_key (required bytes)
         encryption_key, offset = _read_avro_bytes(data, offset)
 
-        # Read aad_prefix (optional: union[null, bytes])
         union_index, offset = _read_avro_long(data, offset)
         if union_index == 0:
             aad_prefix = b""
@@ -94,50 +76,30 @@ def deserialize(data: bytes) -> StandardKeyMetadata:
         else:
             raise ValueError(f"Invalid union index for aad_prefix: {union_index}")
 
-        # Read file_length (optional: union[null, long])
-        file_length = None
+        file_length: int | None = None
         if offset < len(data):
             union_index, offset = _read_avro_long(data, offset)
-            if union_index == 0:
-                file_length = None
-            elif union_index == 1:
+            if union_index == 1:
                 file_length, offset = _read_avro_long(data, offset)
-            else:
+            elif union_index != 0:
                 raise ValueError(f"Invalid union index for file_length: {union_index}")
 
-        return StandardKeyMetadata(
-            encryption_key=encryption_key,
-            aad_prefix=aad_prefix,
-            file_length=file_length,
-        )
+        return StandardKeyMetadata(encryption_key=encryption_key, aad_prefix=aad_prefix, file_length=file_length)
 
     def serialize(self) -> bytes:
-        """Serialize to wire format: ``0x01 version || Avro-encoded fields``."""
-        parts = [bytes([V1])]
-
-        # encryption_key (required bytes)
-        parts.append(_encode_avro_bytes(self.encryption_key))
-
-        # aad_prefix (union[null, bytes])
+        parts = [bytes([V1]), _encode_avro_bytes(self.encryption_key)]
         if self.aad_prefix:
-            parts.append(_encode_avro_long(1))  # union index 1 = bytes
-            parts.append(_encode_avro_bytes(self.aad_prefix))
+            parts += [_encode_avro_long(1), _encode_avro_bytes(self.aad_prefix)]
         else:
-            parts.append(_encode_avro_long(0))  # union index 0 = null
-
-        # file_length (union[null, long])
+            parts.append(_encode_avro_long(0))
         if self.file_length is not None:
-            parts.append(_encode_avro_long(1))  # union index 1 = long
-            parts.append(_encode_avro_long(self.file_length))
+            parts += [_encode_avro_long(1), _encode_avro_long(self.file_length)]
         else:
-            parts.append(_encode_avro_long(0))  # union index 0 = null
-
+            parts.append(_encode_avro_long(0))
         return b"".join(parts)
 
 
 def _encode_avro_long(value: int) -> bytes:
-    """Encode a long as zigzag-encoded Avro varint."""
-    # Zigzag encode
     n = (value << 1) ^ (value >> 63)
     result = bytearray()
     while n & ~0x7F:
@@ -148,5 +110,4 @@ def _encode_avro_long(value: int) -> bytes:
 
 
 def _encode_avro_bytes(data: bytes) -> bytes:
-    """Encode bytes with Avro length prefix."""
     return _encode_avro_long(len(data)) + data

pyiceberg/encryption/kms.py

@@ -14,8 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Key Management Service interfaces and implementations."""
-
 from __future__ import annotations
 
 import importlib
@@ -34,15 +32,11 @@
 
 
 class KeyManagementClient(ABC):
-    """Abstract base class for key management operations."""
-
     @abstractmethod
-    def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes:
-        """Wrap (encrypt) a key using the master key identified by wrapping_key_id."""
+    def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes: ...
 
     @abstractmethod
-    def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes:
-        """Unwrap (decrypt) a wrapped key using the master key identified by wrapping_key_id."""
+    def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes: ...
 
     def initialize(self, properties: dict[str, str]) -> None:  # noqa: B027
         """Initialize the KMS client from catalog/table properties."""
@@ -55,39 +49,26 @@ def __init__(self, master_keys: dict[str, bytes] | None = None) -> None:
         self._master_keys: dict[str, bytes] = dict(master_keys) if master_keys else {}
 
     def initialize(self, properties: dict[str, str]) -> None:
+        prefix = "encryption.kms.key."
         for key, value in properties.items():
-            if key.startswith("encryption.kms.key."):
-                key_id = key[len("encryption.kms.key.") :]
-                self._master_keys[key_id] = bytes.fromhex(value)
+            if key.startswith(prefix):
+                self._master_keys[key[len(prefix) :]] = bytes.fromhex(value)
 
     def wrap_key(self, key: bytes, wrapping_key_id: str) -> bytes:
-        master_key = self._master_keys.get(wrapping_key_id)
-        if master_key is None:
-            raise ValueError(f"Wrapping key not found: {wrapping_key_id}")
-        return aes_gcm_encrypt(master_key, key, aad=None)
+        return aes_gcm_encrypt(self._master(wrapping_key_id), key, aad=None)
 
     def unwrap_key(self, wrapped_key: bytes, wrapping_key_id: str) -> bytes:
+        return aes_gcm_decrypt(self._master(wrapping_key_id), wrapped_key, aad=None)
+
+    def _master(self, wrapping_key_id: str) -> bytes:
         master_key = self._master_keys.get(wrapping_key_id)
         if master_key is None:
             raise ValueError(f"Wrapping key not found: {wrapping_key_id}")
-        return aes_gcm_decrypt(master_key, wrapped_key, aad=None)
+        return master_key
 
 
 def load_kms_client(properties: Properties) -> KeyManagementClient | None:
-    """Load a KMS client from properties using py-kms-impl.
-
-    Follows the same pattern as py-io-impl for FileIO.
-
-    The property 'py-kms-impl' should be a fully qualified Python class name
-    (e.g., 'pyiceberg.encryption.kms.InMemoryKms'). The class must be a
-    subclass of KeyManagementClient.
-
-    Args:
-        properties: Catalog and/or table properties.
-
-    Returns:
-        An initialized KeyManagementClient, or None if py-kms-impl is not set.
-    """
+    """Instantiate a KeyManagementClient from a fully-qualified `py-kms-impl` (or return None)."""
     kms_impl = properties.get(PY_KMS_IMPL)
     if kms_impl is None:
         return None