diff --git a/pyiceberg/encryption/stream.py b/pyiceberg/encryption/stream.py new file mode 100644 index 0000000000..6612fb87ae --- /dev/null +++ b/pyiceberg/encryption/stream.py @@ -0,0 +1,149 @@ +# 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. +"""Format primitives for the AGS1 stream, used to encrypt manifests and manifest lists. + +An AGS1 stream is an 8 byte header followed by a sequence of AES-GCM blocks:: + + "AGS1" || plain_block_size (4 bytes, little endian) + nonce || ciphertext || tag (block 0, up to plain_block_size of plaintext) + nonce || ciphertext || tag (block 1..n, the last of which may be shorter) + +Each block authenticates `aad_prefix || block_index` as additional data, so blocks cannot +be reordered or moved between files. + +The spec gives the last block a non-zero length, which makes a bare header its encoding of an +empty file. Java, iceberg-rust and PyIceberg all require at least one block instead, so an +empty file is a header followed by a single empty block. apache/iceberg#18219 tracks which of +the two forms writers should produce. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from pyiceberg.encryption.ciphers import AesGcmCipher + +_GCM_STREAM_MAGIC = b"AGS1" +_PLAIN_BLOCK_SIZE = 1024 * 1024 +_GCM_STREAM_HEADER_LENGTH = len(_GCM_STREAM_MAGIC) + 4 +_BLOCK_OVERHEAD = AesGcmCipher.NONCE_LENGTH + AesGcmCipher.TAG_LENGTH +_CIPHER_BLOCK_SIZE = _PLAIN_BLOCK_SIZE + _BLOCK_OVERHEAD +_BLOCK_INDEX_LENGTH = 4 +_MAX_BLOCKS = 2 ** (8 * _BLOCK_INDEX_LENGTH) - 1 +_MIN_STREAM_LENGTH = _GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD + + +def _stream_block_aad(aad_prefix: bytes | None, block_index: int) -> bytes: + """Return the additional authenticated data for the block at `block_index`. + + Args: + aad_prefix (bytes | None): The file's AAD prefix, from its key metadata. + block_index (int): The zero-based index of the block within the stream. + """ + return (aad_prefix or b"") + block_index.to_bytes(_BLOCK_INDEX_LENGTH, "little") + + +def _encode_stream_header() -> bytes: + """Encode the AGS1 header that precedes the first block.""" + return _GCM_STREAM_MAGIC + _PLAIN_BLOCK_SIZE.to_bytes(4, "little") + + +def _decode_stream_header(header: bytes) -> int: + """Decode an AGS1 header, returning the plaintext block size it declares. + + Args: + header (bytes): At least `_GCM_STREAM_HEADER_LENGTH` bytes from the start of the stream. + """ + if len(header) < _GCM_STREAM_HEADER_LENGTH: + raise ValueError(f"Invalid AGS1 header: expected {_GCM_STREAM_HEADER_LENGTH} bytes, got {len(header)}") + + if (magic := header[: len(_GCM_STREAM_MAGIC)]) != _GCM_STREAM_MAGIC: + raise ValueError(f"Invalid AGS1 header: magic {magic!r} does not match {_GCM_STREAM_MAGIC!r}") + + plain_block_size = int.from_bytes(header[len(_GCM_STREAM_MAGIC) : _GCM_STREAM_HEADER_LENGTH], "little") + if plain_block_size != _PLAIN_BLOCK_SIZE: + raise ValueError(f"Unsupported AGS1 block size: {plain_block_size} (expected {_PLAIN_BLOCK_SIZE})") + + return plain_block_size + + +@dataclass(frozen=True) +class _Ags1Layout: + """Where each block of an AGS1 stream sits, derived from the trusted encrypted file length. + + Only the final block may hold less than `_PLAIN_BLOCK_SIZE` of plaintext, so the layout + follows from the encrypted length alone, without reading the stream. + """ + + plaintext_length: int + num_blocks: int + last_cipher_block_size: int + + @classmethod + def from_encrypted_length(cls, encrypted_length: int) -> _Ags1Layout: + """Derive the layout of an AGS1 stream that occupies `encrypted_length` bytes. + + Args: + encrypted_length (int): The stream's length, which must be the trusted `file_length` from the file's + `StandardKeyMetadata`, never a file system stat. The spec requires the trusted length because a + stat lets an attacker drop trailing blocks while every remaining block still authenticates. + """ + if encrypted_length < _MIN_STREAM_LENGTH: + raise ValueError(f"Invalid AGS1 stream: expected at least {_MIN_STREAM_LENGTH} bytes, got {encrypted_length}") + + full_blocks, cipher_bytes_in_last_block = divmod(encrypted_length - _GCM_STREAM_HEADER_LENGTH, _CIPHER_BLOCK_SIZE) + if cipher_bytes_in_last_block == 0: + num_blocks, last_cipher_block_size = full_blocks, _CIPHER_BLOCK_SIZE + elif cipher_bytes_in_last_block < _BLOCK_OVERHEAD: + raise ValueError( + f"Truncated AGS1 stream: last block is {cipher_bytes_in_last_block} bytes, expected at least {_BLOCK_OVERHEAD}" + ) + else: + num_blocks, last_cipher_block_size = full_blocks + 1, cipher_bytes_in_last_block + + if num_blocks > _MAX_BLOCKS: + raise ValueError(f"AGS1 streams hold at most {_MAX_BLOCKS} blocks, but {encrypted_length} bytes needs {num_blocks}") + + return cls( + plaintext_length=(num_blocks - 1) * _PLAIN_BLOCK_SIZE + last_cipher_block_size - _BLOCK_OVERHEAD, + num_blocks=num_blocks, + last_cipher_block_size=last_cipher_block_size, + ) + + def _check_block_index(self, block_index: int) -> None: + if not 0 <= block_index < self.num_blocks: + raise ValueError(f"Block index out of range: {block_index} (stream holds {self.num_blocks} blocks)") + + def cipher_block_size(self, block_index: int) -> int: + """Return the encrypted size of the block at `block_index`.""" + self._check_block_index(block_index) + return self.last_cipher_block_size if block_index == self.num_blocks - 1 else _CIPHER_BLOCK_SIZE + + def plain_block_size(self, block_index: int) -> int: + """Return the plaintext size of the block at `block_index`.""" + return self.cipher_block_size(block_index) - _BLOCK_OVERHEAD + + def encrypted_block_offset(self, block_index: int) -> int: + """Return the offset of the block at `block_index` within the encrypted stream.""" + self._check_block_index(block_index) + return _GCM_STREAM_HEADER_LENGTH + block_index * _CIPHER_BLOCK_SIZE + + def block_index_for(self, plaintext_offset: int) -> int: + """Return the index of the block holding `plaintext_offset`.""" + if not 0 <= plaintext_offset < self.plaintext_length: + raise ValueError(f"Plaintext offset out of range: {plaintext_offset} (stream holds {self.plaintext_length} bytes)") + return plaintext_offset // _PLAIN_BLOCK_SIZE diff --git a/tests/encryption/ags1/GenerateAgs1Fixtures.java b/tests/encryption/ags1/GenerateAgs1Fixtures.java new file mode 100644 index 0000000000..a8a2dc9e82 --- /dev/null +++ b/tests/encryption/ags1/GenerateAgs1Fixtures.java @@ -0,0 +1,69 @@ +/* + * 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 java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import org.apache.iceberg.Files; +import org.apache.iceberg.encryption.AesGcmOutputFile; +import org.apache.iceberg.io.PositionOutputStream; + +/** Writes the AGS1 fixtures in this directory with Java's AesGcmOutputStream. See README.md. */ +public class GenerateAgs1Fixtures { + static final int PLAIN_BLOCK_SIZE = 1024 * 1024; + static final byte[] KEY = new byte[16]; + static final byte[] AAD_PREFIX = "pyiceberg-ags1".getBytes(StandardCharsets.UTF_8); + + static { + for (int i = 0; i < KEY.length; i++) { + KEY[i] = (byte) i; + } + } + + /** Byte i is i % 251. The prime period shifts phase across every 1 MiB block boundary. */ + static byte[] plaintext(int length) { + byte[] out = new byte[length]; + for (int i = 0; i < length; i++) { + out[i] = (byte) (i % 251); + } + return out; + } + + static void write(File dir, String name, int length, byte[] aadPrefix) throws IOException { + File target = new File(dir, name); + if (target.exists() && !target.delete()) { + throw new IOException("Could not delete " + target); + } + + AesGcmOutputFile encrypted = new AesGcmOutputFile(Files.localOutput(target), KEY, aadPrefix); + try (PositionOutputStream stream = encrypted.create()) { + stream.write(plaintext(length)); + } + + System.out.printf("%-26s plaintext=%-8d encrypted=%d%n", name, length, target.length()); + } + + public static void main(String[] args) throws IOException { + File dir = new File(args.length > 0 ? args[0] : "."); + write(dir, "empty.ags1", 0, AAD_PREFIX); + write(dir, "partial-block.ags1", 100, AAD_PREFIX); + write(dir, "partial-block-no-aad.ags1", 100, null); + write(dir, "aligned-multi-block.ags1", 2 * PLAIN_BLOCK_SIZE, AAD_PREFIX); + } +} diff --git a/tests/encryption/ags1/README.md b/tests/encryption/ags1/README.md new file mode 100644 index 0000000000..30d5c907f9 --- /dev/null +++ b/tests/encryption/ags1/README.md @@ -0,0 +1,77 @@ + + +# AGS1 cross-client test fixtures + +These files were written by Java's `AesGcmOutputStream` (Apache Iceberg 1.11.0), not +by PyIceberg. They exist so PyIceberg's AGS1 support is checked against another +implementation's bytes rather than only against its own round trip. All four also +decrypt with iceberg-rust 0.10.1, through its `EncryptedInputFile`. + +## Fixtures + +| File | Encrypted size | Plaintext | Pins | +| --- | --- | --- | --- | +| `empty.ags1` | 36 B | 0 B | Java writes an 8 byte header **plus one empty block** for an empty file, not a bare header | +| `partial-block.ags1` | 136 B | 100 B | Header, nonce/tag layout, and a single short block | +| `partial-block-no-aad.ags1` | 136 B | 100 B | The same stream with a null AAD prefix, so the block index alone is the AAD | +| `aligned-multi-block.ags1` | 2097216 B | 2 MiB | Two full blocks: the little-endian block index in each block's AAD, and that a block-aligned write appends **no** trailing empty block | + +`empty.ags1` is worth calling out. The spec says the last block has a non-zero +length, which makes a bare 8 byte header its encoding of an empty file. Java instead +writes 36 bytes and its `AesGcmInputFile` rejects anything shorter, as does iceberg-rust +since apache/iceberg-rust#3236. PyIceberg follows the implementations rather than the +spec here: `_MIN_STREAM_LENGTH` is 36, so a bare header is rejected rather than read as +an empty stream. apache/iceberg#18219 tracks which of the two forms writers should +produce. + +Block-aligned and partial *single* block variants are deliberately not checked in. +The 1 MiB block size is hard-coded, so each would add another 1 MiB of +incompressible ciphertext without covering a case the four files above miss. + +## Parameters + +Every fixture uses: + +- **Key**: 16 bytes, `0x00` through `0x0f` +- **AAD prefix**: ASCII `pyiceberg-ags1`, except `partial-block-no-aad.ags1`, which has none +- **Plaintext**: byte `i` is `i % 251`. The period is prime and therefore coprime with the + 1 MiB block size, so the pattern shifts phase at every block boundary and a + misordered or misindexed block is detectable from the plaintext alone + +## Regenerating + +Each block uses a fresh random nonce, so regenerating produces different bytes. +The file lengths, and the plaintext each file decrypts to, are deterministic. Tests +decrypt these fixtures rather than comparing them byte for byte, so a regeneration +is safe as long as the parameters above are unchanged. + +From this directory, with a JDK 17 or later: + + +```bash +V=1.11.0 +for a in iceberg-core iceberg-api iceberg-bundled-guava; do + curl -sfLO "https://repo1.maven.org/maven2/org/apache/iceberg/$a/$V/$a-$V.jar" +done +java -cp "iceberg-api-$V.jar:iceberg-bundled-guava-$V.jar:iceberg-core-$V.jar" \ + GenerateAgs1Fixtures.java . +rm iceberg-*-$V.jar +``` + diff --git a/tests/encryption/ags1/aligned-multi-block.ags1 b/tests/encryption/ags1/aligned-multi-block.ags1 new file mode 100644 index 0000000000..9db2afe423 Binary files /dev/null and b/tests/encryption/ags1/aligned-multi-block.ags1 differ diff --git a/tests/encryption/ags1/empty.ags1 b/tests/encryption/ags1/empty.ags1 new file mode 100644 index 0000000000..2482581657 Binary files /dev/null and b/tests/encryption/ags1/empty.ags1 differ diff --git a/tests/encryption/ags1/partial-block-no-aad.ags1 b/tests/encryption/ags1/partial-block-no-aad.ags1 new file mode 100644 index 0000000000..a938a0df7e Binary files /dev/null and b/tests/encryption/ags1/partial-block-no-aad.ags1 differ diff --git a/tests/encryption/ags1/partial-block.ags1 b/tests/encryption/ags1/partial-block.ags1 new file mode 100644 index 0000000000..ea9fb4b2eb Binary files /dev/null and b/tests/encryption/ags1/partial-block.ags1 differ diff --git a/tests/encryption/test_stream.py b/tests/encryption/test_stream.py new file mode 100644 index 0000000000..799e241a10 --- /dev/null +++ b/tests/encryption/test_stream.py @@ -0,0 +1,306 @@ +# 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. + +from pathlib import Path + +import pytest + +from pyiceberg.encryption.ciphers import AesGcmCipher, SecureKey +from pyiceberg.encryption.stream import ( + _BLOCK_OVERHEAD, + _CIPHER_BLOCK_SIZE, + _GCM_STREAM_HEADER_LENGTH, + _GCM_STREAM_MAGIC, + _MAX_BLOCKS, + _MIN_STREAM_LENGTH, + _PLAIN_BLOCK_SIZE, + _Ags1Layout, + _decode_stream_header, + _encode_stream_header, + _stream_block_aad, +) + +KEY = SecureKey(b"0123456789012345") +AAD_PREFIX = b"0123456789abcdef" + +# The header a Java `AesGcmOutputStream` writes: "AGS1" then 1 MiB as a little-endian int32. +JAVA_HEADER = b"AGS1\x00\x00\x10\x00" + +# Streams written by Java's `AesGcmOutputStream`, with the parameters documented in ags1/README.md. +AGS1_FIXTURES = Path(__file__).parent / "ags1" +FIXTURE_KEY = SecureKey(bytes(range(16))) +FIXTURE_AAD_PREFIX = b"pyiceberg-ags1" + + +def build_stream(plaintext: bytes, aad_prefix: bytes | None = AAD_PREFIX) -> bytes: + """Encrypt `plaintext` into an AGS1 stream, as an output stream implementation would. + + An empty plaintext still gets one empty block, which is the form Java writes and the reader accepts. + """ + blocks = [ + AesGcmCipher(KEY).encrypt(plaintext[start : start + _PLAIN_BLOCK_SIZE], _stream_block_aad(aad_prefix, index)) + for index, start in enumerate(range(0, len(plaintext), _PLAIN_BLOCK_SIZE) or [0]) + ] + return _encode_stream_header() + b"".join(blocks) + + +def decrypt_stream(stream: bytes, key: SecureKey, aad_prefix: bytes | None) -> bytes: + """Decrypt an AGS1 stream block by block, as an input stream implementation would.""" + layout = _Ags1Layout.from_encrypted_length(len(stream)) + return b"".join( + AesGcmCipher(key).decrypt( + stream[layout.encrypted_block_offset(index) : layout.encrypted_block_offset(index) + layout.cipher_block_size(index)], + _stream_block_aad(aad_prefix, index), + ) + for index in range(layout.num_blocks) + ) + + +def fixture_plaintext(length: int) -> bytes: + """Return the plaintext the AGS1 fixtures encrypt: byte `i` is `i % 251`.""" + return (bytes(range(251)) * (length // 251 + 1))[:length] + + +def test_format_constants() -> None: + assert _GCM_STREAM_MAGIC == b"AGS1" + assert _PLAIN_BLOCK_SIZE == 1024 * 1024 + assert _GCM_STREAM_HEADER_LENGTH == 8 + assert _BLOCK_OVERHEAD == 28 + assert _CIPHER_BLOCK_SIZE == _PLAIN_BLOCK_SIZE + _BLOCK_OVERHEAD + assert _MAX_BLOCKS == 2**32 - 1 + assert _MIN_STREAM_LENGTH == 36 + + +def test_encode_stream_header_matches_java() -> None: + assert _encode_stream_header() == JAVA_HEADER + + +def test_decode_stream_header() -> None: + assert _decode_stream_header(JAVA_HEADER) == _PLAIN_BLOCK_SIZE + assert _decode_stream_header(_encode_stream_header()) == _PLAIN_BLOCK_SIZE + + +def test_decode_stream_header_ignores_trailing_block_bytes() -> None: + assert _decode_stream_header(JAVA_HEADER + b"block bytes") == _PLAIN_BLOCK_SIZE + + +@pytest.mark.parametrize("length", [0, 4, 7]) +def test_decode_stream_header_rejects_a_short_header(length: int) -> None: + with pytest.raises(ValueError, match=f"Invalid AGS1 header: expected 8 bytes, got {length}"): + _decode_stream_header(bytes(length)) + + +def test_decode_stream_header_rejects_the_wrong_magic() -> None: + with pytest.raises(ValueError, match="magic b'AGS2' does not match b'AGS1'"): + _decode_stream_header(b"AGS2\x00\x00\x10\x00") + + +def test_decode_stream_header_rejects_an_unsupported_block_size() -> None: + with pytest.raises(ValueError, match=f"Unsupported AGS1 block size: 512 \\(expected {_PLAIN_BLOCK_SIZE}\\)"): + _decode_stream_header(_GCM_STREAM_MAGIC + (512).to_bytes(4, "little")) + + +@pytest.mark.parametrize( + "block_index, expected", + [(0, b"\x00\x00\x00\x00"), (1, b"\x01\x00\x00\x00"), (258, b"\x02\x01\x00\x00"), (_MAX_BLOCKS, b"\xff\xff\xff\xff")], +) +def test_stream_block_aad_encodes_the_index_little_endian(block_index: int, expected: bytes) -> None: + assert _stream_block_aad(None, block_index) == expected + assert _stream_block_aad(b"", block_index) == expected + assert _stream_block_aad(AAD_PREFIX, block_index) == AAD_PREFIX + expected + + +@pytest.mark.parametrize( + "encrypted_length, plaintext_length, num_blocks, last_cipher_block_size", + [ + (_GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD, 0, 1, _BLOCK_OVERHEAD), + (_GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD + 100, 100, 1, _BLOCK_OVERHEAD + 100), + (_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE, _PLAIN_BLOCK_SIZE, 1, _CIPHER_BLOCK_SIZE), + (_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE + _BLOCK_OVERHEAD + 5, _PLAIN_BLOCK_SIZE + 5, 2, _BLOCK_OVERHEAD + 5), + (_GCM_STREAM_HEADER_LENGTH + 2 * _CIPHER_BLOCK_SIZE, 2 * _PLAIN_BLOCK_SIZE, 2, _CIPHER_BLOCK_SIZE), + ], +) +def test_layout_from_encrypted_length( + encrypted_length: int, plaintext_length: int, num_blocks: int, last_cipher_block_size: int +) -> None: + layout = _Ags1Layout.from_encrypted_length(encrypted_length) + + assert layout == _Ags1Layout( + plaintext_length=plaintext_length, num_blocks=num_blocks, last_cipher_block_size=last_cipher_block_size + ) + + +@pytest.mark.parametrize("encrypted_length", [0, 1, 7, _GCM_STREAM_HEADER_LENGTH, _MIN_STREAM_LENGTH - 1]) +def test_layout_rejects_a_stream_shorter_than_one_block(encrypted_length: int) -> None: + """A header alone is not a stream, matching Java's `MIN_STREAM_LENGTH` and iceberg-rust.""" + with pytest.raises(ValueError, match=f"expected at least {_MIN_STREAM_LENGTH} bytes, got {encrypted_length}"): + _Ags1Layout.from_encrypted_length(encrypted_length) + + +@pytest.mark.parametrize("last_block_size", [1, 27]) +def test_layout_rejects_a_truncated_last_block(last_block_size: int) -> None: + with pytest.raises(ValueError, match=f"last block is {last_block_size} bytes, expected at least 28"): + _Ags1Layout.from_encrypted_length(_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE + last_block_size) + + +def test_layout_rejects_more_blocks_than_the_index_can_address() -> None: + encrypted_length = _GCM_STREAM_HEADER_LENGTH + (_MAX_BLOCKS + 1) * _CIPHER_BLOCK_SIZE + + with pytest.raises(ValueError, match=f"AGS1 streams hold at most {_MAX_BLOCKS} blocks"): + _Ags1Layout.from_encrypted_length(encrypted_length) + + +def test_layout_block_sizes_and_offsets() -> None: + layout = _Ags1Layout.from_encrypted_length(_GCM_STREAM_HEADER_LENGTH + 2 * _CIPHER_BLOCK_SIZE + _BLOCK_OVERHEAD + 7) + + assert layout.num_blocks == 3 + assert layout.cipher_block_size(0) == layout.cipher_block_size(1) == _CIPHER_BLOCK_SIZE + assert layout.plain_block_size(0) == layout.plain_block_size(1) == _PLAIN_BLOCK_SIZE + assert layout.cipher_block_size(2) == _BLOCK_OVERHEAD + 7 + assert layout.plain_block_size(2) == 7 + assert layout.encrypted_block_offset(0) == _GCM_STREAM_HEADER_LENGTH + assert layout.encrypted_block_offset(1) == _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE + assert layout.encrypted_block_offset(2) == _GCM_STREAM_HEADER_LENGTH + 2 * _CIPHER_BLOCK_SIZE + + +@pytest.mark.parametrize("block_index", [-1, 1, 2]) +def test_layout_rejects_an_out_of_range_block_index(block_index: int) -> None: + layout = _Ags1Layout.from_encrypted_length(_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE) + + with pytest.raises(ValueError, match=f"Block index out of range: {block_index} \\(stream holds 1 blocks\\)"): + layout.cipher_block_size(block_index) + + with pytest.raises(ValueError, match=f"Block index out of range: {block_index}"): + layout.encrypted_block_offset(block_index) + + +@pytest.mark.parametrize( + "plaintext_offset, expected", + [(0, 0), (1, 0), (_PLAIN_BLOCK_SIZE - 1, 0), (_PLAIN_BLOCK_SIZE, 1), (_PLAIN_BLOCK_SIZE + 6, 1)], +) +def test_layout_block_index_for_plaintext_offset(plaintext_offset: int, expected: int) -> None: + layout = _Ags1Layout.from_encrypted_length(_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE + _BLOCK_OVERHEAD + 7) + + assert layout.block_index_for(plaintext_offset) == expected + + +@pytest.mark.parametrize("plaintext_offset", [-1, _PLAIN_BLOCK_SIZE]) +def test_layout_rejects_an_out_of_range_plaintext_offset(plaintext_offset: int) -> None: + layout = _Ags1Layout.from_encrypted_length(_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE) + + with pytest.raises(ValueError, match=f"Plaintext offset out of range: {plaintext_offset}"): + layout.block_index_for(plaintext_offset) + + +@pytest.mark.parametrize("plaintext_length", [0, 1, 100, _PLAIN_BLOCK_SIZE, _PLAIN_BLOCK_SIZE + 7, 2 * _PLAIN_BLOCK_SIZE]) +def test_layout_describes_a_real_stream(plaintext_length: int) -> None: + """The layout derived from a stream's length must match the stream that was written.""" + plaintext = fixture_plaintext(plaintext_length) + stream = build_stream(plaintext) + + layout = _Ags1Layout.from_encrypted_length(len(stream)) + + assert _decode_stream_header(stream) == _PLAIN_BLOCK_SIZE + assert layout.plaintext_length == plaintext_length + assert layout.num_blocks == max(1, -(-plaintext_length // _PLAIN_BLOCK_SIZE)) + + decrypted = b"" + for index in range(layout.num_blocks): + offset = layout.encrypted_block_offset(index) + block = stream[offset : offset + layout.cipher_block_size(index)] + decrypted += AesGcmCipher(KEY).decrypt(block, _stream_block_aad(AAD_PREFIX, index)) + + assert decrypted == plaintext + + +def test_blocks_cannot_be_reordered() -> None: + stream = build_stream(bytes(_PLAIN_BLOCK_SIZE + 7)) + layout = _Ags1Layout.from_encrypted_length(len(stream)) + first_block = stream[layout.encrypted_block_offset(0) : layout.encrypted_block_offset(1)] + + with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): + AesGcmCipher(KEY).decrypt(first_block, _stream_block_aad(AAD_PREFIX, 1)) + + +def test_blocks_cannot_be_moved_between_files() -> None: + stream = build_stream(bytes(100)) + layout = _Ags1Layout.from_encrypted_length(len(stream)) + block = stream[layout.encrypted_block_offset(0) :] + + with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): + AesGcmCipher(KEY).decrypt(block, _stream_block_aad(b"another file's prefix", 0)) + + +@pytest.mark.parametrize( + "name, plaintext_length, num_blocks, aad_prefix", + [ + ("empty.ags1", 0, 1, FIXTURE_AAD_PREFIX), + ("partial-block.ags1", 100, 1, FIXTURE_AAD_PREFIX), + ("partial-block-no-aad.ags1", 100, 1, None), + ("aligned-multi-block.ags1", 2 * _PLAIN_BLOCK_SIZE, 2, FIXTURE_AAD_PREFIX), + ], +) +def test_decrypts_a_java_written_stream(name: str, plaintext_length: int, num_blocks: int, aad_prefix: bytes | None) -> None: + """A stream written by Java must decrypt with the layout derived from its length alone.""" + stream = (AGS1_FIXTURES / name).read_bytes() + + layout = _Ags1Layout.from_encrypted_length(len(stream)) + + assert stream[:_GCM_STREAM_HEADER_LENGTH] == _encode_stream_header() + assert _decode_stream_header(stream) == _PLAIN_BLOCK_SIZE + assert layout.plaintext_length == plaintext_length + assert layout.num_blocks == num_blocks + assert decrypt_stream(stream, FIXTURE_KEY, aad_prefix) == fixture_plaintext(plaintext_length) + + +def test_java_encodes_an_empty_file_as_one_empty_block() -> None: + """Java writes a header plus one empty block for an empty file, which is the shortest stream accepted.""" + stream = (AGS1_FIXTURES / "empty.ags1").read_bytes() + + assert len(stream) == _MIN_STREAM_LENGTH + assert _Ags1Layout.from_encrypted_length(len(stream)) == _Ags1Layout( + plaintext_length=0, num_blocks=1, last_cipher_block_size=_BLOCK_OVERHEAD + ) + + +def test_java_appends_no_trailing_block_to_a_block_aligned_stream() -> None: + """A block-aligned write ends on its last full block, so the length holds no extra empty block.""" + stream = (AGS1_FIXTURES / "aligned-multi-block.ags1").read_bytes() + + assert len(stream) == _GCM_STREAM_HEADER_LENGTH + 2 * _CIPHER_BLOCK_SIZE + + +def test_a_truncated_java_stream_still_authenticates() -> None: + """Dropping a trailing block leaves every remaining block valid, so the length must come from key metadata.""" + stream = (AGS1_FIXTURES / "aligned-multi-block.ags1").read_bytes() + + truncated = stream[: _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE] + + assert decrypt_stream(truncated, FIXTURE_KEY, FIXTURE_AAD_PREFIX) == fixture_plaintext(_PLAIN_BLOCK_SIZE) + + +def test_the_trusted_length_exposes_a_truncated_stream() -> None: + """The layout from the trusted length outruns a truncated file, which is how a reader detects the drop.""" + stream = (AGS1_FIXTURES / "aligned-multi-block.ags1").read_bytes() + truncated = stream[: _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE] + + layout = _Ags1Layout.from_encrypted_length(len(stream)) + last_block = layout.num_blocks - 1 + + assert layout.plaintext_length == 2 * _PLAIN_BLOCK_SIZE + assert layout.encrypted_block_offset(last_block) + layout.cipher_block_size(last_block) > len(truncated) + assert _Ags1Layout.from_encrypted_length(len(truncated)).plaintext_length == _PLAIN_BLOCK_SIZE