diff --git a/dissect/util/compression/__init__.py b/dissect/util/compression/__init__.py index 549eb46..aae667e 100644 --- a/dissect/util/compression/__init__.py +++ b/dissect/util/compression/__init__.py @@ -1,3 +1,12 @@ +"""Compression algorithms: LZ4, LZO, LZNT1, LZXPRESS, DEFLATE, ZLIB, and others. + +Selects between native Rust implementations of LZ4 and LZO (when available +via dissect.util._native) and pure-Python fallbacks. The native versions are +exposed as ``lz4_native`` / ``lzo_native`` and the Python versions as +``lz4_python`` / ``lzo_python``; importing ``lz4`` or ``lzo`` gives whichever +is available, preferring native. +""" + from __future__ import annotations from dissect.util.compression import lz4, lzo @@ -21,7 +30,7 @@ # dissect.util.compression.lzo_python. # # Note that the pure Python implementation and the Rust implementation are NOT a full replacement -# for the "official" lz4 and lzo Python packages: only the decompress() function is implemented. +# for the "official" lz4 and lzo Python packages. try: from dissect.util import _native @@ -31,6 +40,7 @@ lz4_native = lzo_native = None __all__ = [ + "deflate", "lz4", "lz4_native", "lz4_python", @@ -42,6 +52,8 @@ "lzo_python", "lzvn", "lzxpress", + "lzxpress9", "lzxpress_huffman", "sevenbit", + "zlibstream", ] diff --git a/dissect/util/compression/deflate.py b/dissect/util/compression/deflate.py new file mode 100644 index 0000000..91b9ce6 --- /dev/null +++ b/dissect/util/compression/deflate.py @@ -0,0 +1,53 @@ +"""Raw DEFLATE (RFC 1951) compression and decompression without zlib or gzip wrappers. + +Used by ntdll.dll RtlCompressBuffer with COMPRESSION_FORMAT_DEFLATE (0x0007). +""" + +from __future__ import annotations + +import zlib as _zlib +from typing import BinaryIO + + +def decompress(src: bytes | BinaryIO, max_size: int | None = None) -> bytes: + """Decompress a raw DEFLATE stream (RFC 1951, no zlib/gzip wrapper). + + Args: + src: File-like object or bytes to decompress. + max_size: Optional maximum output size in bytes. When set, raises + a zlib.error if the decompressed output would exceed this limit. + + Returns: + The decompressed data. + + Raises: + zlib.error: If the input is malformed or the output exceeds ``max_size``. + """ + if hasattr(src, "read"): + src = src.read() + if max_size is not None: + dc = _zlib.decompressobj(wbits=-15) + result = dc.decompress(src, max_size) + if dc.unconsumed_tail: + raise _zlib.error("DEFLATE output exceeds max_size") + return result + return _zlib.decompress(src, -15) + + +def compress(src: bytes | BinaryIO, level: int = 1) -> bytes: + """Compress data into a raw DEFLATE stream (RFC 1951, no wrapper). + + Args: + src: File-like object or bytes to compress. + level: Compression level 0-9 (default 1). + + Returns: + The compressed raw DEFLATE bitstream. + + Raises: + zlib.error: If the input cannot be compressed or the level is invalid. + """ + if hasattr(src, "read"): + src = src.read() + c = _zlib.compressobj(level, _zlib.DEFLATED, -15) + return c.compress(src) + c.flush() diff --git a/dissect/util/compression/lz4.py b/dissect/util/compression/lz4.py index 6736bd0..da7dc4c 100644 --- a/dissect/util/compression/lz4.py +++ b/dissect/util/compression/lz4.py @@ -1,3 +1,5 @@ +"""LZ4 block-format compression and decompression.""" + from __future__ import annotations import io @@ -38,6 +40,9 @@ def decompress( Returns: The decompressed data. + + Raises: + CorruptDataError: If the input is truncated, malformed, or exceeds ``uncompressed_size``. """ if not hasattr(src, "read"): src = io.BytesIO(src) @@ -93,3 +98,115 @@ def decompress( dst = bytes(dst) return dst[:uncompressed_size] if uncompressed_size > 0 else dst + + +# --- LZ4 block encoder --- + +_MIN_MATCH = 4 +"""Minimum match length (lz4_Block_format.md).""" + +_MAX_OFFSET = 0xFFFF +"""Largest encodable back-reference distance (lz4_Block_format.md).""" + +_TOKEN_MAX = 15 +"""Nibble value meaning "length continues in extra bytes" (lz4_Block_format.md).""" + +_LAST_LITERALS = 5 +"""End-of-block invariant: the last 5 bytes are always literals (lz4_Block_format.md).""" + +_MF_LIMIT = 12 +"""End-of-block invariant: last match starts at least 12 bytes before end (lz4_Block_format.md).""" + + +def _emit_length_extension(out: bytearray, remainder: int) -> None: + """Emit 0xFF-continuation bytes for a length whose nibble was 15. + + Per lz4_Block_format.md each byte (0-255) adds to the length, and a 255 + byte forces another. An exact multiple of 255 is terminated by an + explicit 0 byte. + """ + while remainder >= 0xFF: + out.append(0xFF) + remainder -= 0xFF + out.append(remainder) + + +def _emit_sequence(out: bytearray, literals: bytes, offset: int, match_len: int) -> None: + """Emit one full sequence: token, literal run, offset, match length. + + A ``match_len`` of 0 marks the final literals-only sequence (no offset + field) -- the block-format end condition that the last sequence contains + only literals. + """ + literal_len = len(literals) + literal_nibble = min(literal_len, _TOKEN_MAX) + match_nibble = 0 if match_len == 0 else min(match_len - _MIN_MATCH, _TOKEN_MAX) + out.append((literal_nibble << 4) | match_nibble) + if literal_nibble == _TOKEN_MAX: + _emit_length_extension(out, literal_len - _TOKEN_MAX) + out += literals + if match_len == 0: + return + out += offset.to_bytes(2, "little") + if match_nibble == _TOKEN_MAX: + _emit_length_extension(out, match_len - _MIN_MATCH - _TOKEN_MAX) + + +def compress(src: bytes | BinaryIO, return_bytearray: bool = False) -> bytes: + """LZ4 compress to a headerless block. + + Greedy hash-table matcher honouring all end-of-block invariants from the + LZ4 block format specification: matches are at least 4 bytes, reach back + at most 65535 bytes, start at least 12 bytes before the end, and stop 5 + bytes short so the block ends in a literals-only sequence. + + Args: + src: Data to compress, as bytes or a file-like object. + return_bytearray: Whether to return ``bytearray`` instead of ``bytes``. + + Returns: + The compressed LZ4 block (no frame header, no checksum). + """ + if hasattr(src, "read"): + src = src.read() + + data = bytes(src) + n = len(data) + out = bytearray() + + if n == 0: + # A single zero token: empty final literal run. + out.append(0x00) + if not return_bytearray: + return bytes(out) + return out + + table: dict[bytes, int] = {} + match_limit = n - _LAST_LITERALS + anchor = 0 + pos = 0 + + while pos <= n - _MF_LIMIT: + key = data[pos : pos + _MIN_MATCH] + candidate = table.get(key) + table[key] = pos + + if candidate is None or pos - candidate > _MAX_OFFSET: + pos += 1 + continue + + # dict key equality guarantees the first _MIN_MATCH bytes match. + match_len = _MIN_MATCH + while pos + match_len < match_limit and data[candidate + match_len] == data[pos + match_len]: + match_len += 1 + + _emit_sequence(out, data[anchor:pos], pos - candidate, match_len) + pos += match_len + anchor = pos + + # Final literals-only sequence. + _emit_sequence(out, data[anchor:], 0, 0) + + if not return_bytearray: + return bytes(out) + return out diff --git a/dissect/util/compression/lznt1.py b/dissect/util/compression/lznt1.py index 5981b86..acdecaf 100644 --- a/dissect/util/compression/lznt1.py +++ b/dissect/util/compression/lznt1.py @@ -1,3 +1,5 @@ +"""LZNT1 compression and decompression per [MS-XCA] S2.5.""" + # Reference: https://github.com/google/rekall/blob/master/rekall-core/rekall/plugins/filesystems/lznt1.py from __future__ import annotations @@ -33,6 +35,9 @@ def decompress(src: bytes | BinaryIO) -> bytes: Returns: The decompressed data. + + Raises: + struct.error: If the input is truncated or malformed. """ if not hasattr(src, "read"): src = io.BytesIO(src) @@ -90,3 +95,138 @@ def decompress(src: bytes | BinaryIO) -> bytes: dst.write(data) return dst.getvalue() + + +# --- Compression helpers ([MS-XCA] S2.5.2-2.5.4) --- + +_CHUNK_SIZE = 4096 +"""Uncompressed bytes consumed per chunk; [MS-XCA] S2.5.3 mandates 4096-byte units.""" + +_MIN_MATCH = 3 +"""Minimum match length; stored length is actual minus 3 ([MS-XCA] S2.5.1.4).""" + +_MAX_MATCH_CANDIDATES = 256 +"""Encoder-only cap on hash-chain probes per position; bounds worst-case time.""" + + +def _displacement_bits(pos: int) -> int: + """Return the displacement bit width at a chunk position per [MS-XCA] section 2.5.1.4. + + Reuses the module-level DISPLACEMENT_TABLE, adding the minimum D of 4 that the + table's values are relative to. + """ + if pos == 0: + return 4 + return DISPLACEMENT_TABLE[pos - 1] + 4 + + +def _find_match(chunk: bytes, pos: int, table: dict[bytes, list[int]]) -> tuple[int, int]: + """Return (length, displacement) of the best match at pos, or (0, 0) if none. + + Searches the hash chain in reverse (most recent first) for the longest match + within the displacement window allowed by the current chunk position. + """ + if len(chunk) - pos < _MIN_MATCH: + return 0, 0 + + disp_bits = _displacement_bits(pos) + max_displacement = min(pos, 1 << disp_bits) + max_length = min(len(chunk) - pos, (1 << (16 - disp_bits)) - 1 + _MIN_MATCH) + + best_length = 0 + best_displacement = 0 + + candidates = table.get(chunk[pos : pos + _MIN_MATCH], []) + for candidate in reversed(candidates[-_MAX_MATCH_CANDIDATES:]): + if pos - candidate > max_displacement: + break + length = _MIN_MATCH + while length < max_length and chunk[candidate + length] == chunk[pos + length]: + length += 1 + if length > best_length: + best_length = length + best_displacement = pos - candidate + if length == max_length: + break + + return best_length, best_displacement + + +def _compress_chunk(chunk: bytes) -> bytes: + """Encode one chunk of plaintext into an LZNT1 flag-group body. + + Each flag group is a flag byte followed by up to eight data elements. A clear + flag bit emits a literal byte; a set bit emits a 16-bit compressed word encoding + a (displacement, length) back-reference per [MS-XCA] section 2.5.1.3. + """ + body = bytearray() + table: dict[bytes, list[int]] = {} + group_flags = 0 + group_items = bytearray() + group_count = 0 + pos = 0 + + while pos < len(chunk): + length, displacement = _find_match(chunk, pos, table) + + if length >= _MIN_MATCH: + disp_bits = _displacement_bits(pos) + word = (displacement - 1) << (16 - disp_bits) | (length - _MIN_MATCH) + group_items += struct.pack(" bytes: + """LZNT1 compress from a file-like object or bytes. + + Consumes the input in 4096-byte units per [MS-XCA] section 2.5.3, emitting one + chunk each: a compressed chunk when the flag-group body is smaller than the raw + data, otherwise an uncompressed chunk. + + Args: + src: File-like object or bytes to compress. + + Returns: + The compressed data. + """ + if hasattr(src, "read"): + src = src.read() + + result = bytearray() + + for start in range(0, len(src), _CHUNK_SIZE): + chunk = src[start : start + _CHUNK_SIZE] + body = _compress_chunk(chunk) + + if len(body) < len(chunk): + header = COMPRESSED_MASK | SIGNATURE_MASK | (len(body) + 2 - 3) + result += struct.pack(" bytes: + +def decompress(src: bytes | BinaryIO, max_size: int | None = None) -> bytes: """LZXPRESS decompress from a file-like object or bytes. Args: src: File-like object or bytes to decompress. + max_size: Optional maximum output size in bytes. When set, raises + ValueError if the decompressed output would exceed this limit. Returns: The decompressed data. + + Raises: + ValueError: If the output exceeds ``max_size`` or the stream contains an invalid match length. """ if not hasattr(src, "read"): src = io.BytesIO(src) @@ -36,6 +56,8 @@ def decompress(src: bytes | BinaryIO) -> bytes: buffered_flags_count -= 1 if buffered_flags & (1 << buffered_flags_count) == 0: + if max_size is not None and len(dst) >= max_size: + raise ValueError(f"decompressed output exceeds max_size of {max_size}") dst.append(ord(src.read(1))) else: if src.tell() - offset == size: @@ -71,6 +93,9 @@ def decompress(src: bytes | BinaryIO) -> bytes: match_length += 7 match_length += 3 + if max_size is not None and len(dst) + match_length > max_size: + raise ValueError(f"decompressed output exceeds max_size of {max_size}") + remaining = match_length while remaining > 0: match_size = min(remaining, match_offset) @@ -78,3 +103,152 @@ def decompress(src: bytes | BinaryIO) -> bytes: remaining -= match_size return bytes(dst) + + +# --- Compression ([MS-XCA] S2.3.4) --- + + +class _Encoder: + """Mutable output buffer for the [MS-XCA] S2.3.4 encoder.""" + + def __init__(self) -> None: + self.out = bytearray(4) + self.flags = 0 + self.flag_count = 0 + self.flag_pos = 0 + self.nibble_pos = -1 + + def _push_flag(self, bit: int) -> None: + """Append one literal/match flag bit; flush the word every 32 bits.""" + self.flags = self.flags << 1 | bit + self.flag_count += 1 + if self.flag_count == _FLAG_BITS: + self.out[self.flag_pos : self.flag_pos + 4] = self.flags.to_bytes(4, "little") + self.flags = 0 + self.flag_count = 0 + self.flag_pos = len(self.out) + self.out += b"\x00\x00\x00\x00" + + def emit_literal(self, byte: int) -> None: + """Copy one raw byte to the output under a 0 flag bit.""" + self.out.append(byte) + self._push_flag(0) + + def _emit_long_length(self, extra: int) -> None: + """Encode length - 3 - 7 through the nibble/byte/word/dword ladder.""" + if self.nibble_pos < 0: + self.nibble_pos = len(self.out) + self.out.append(min(extra, _NIBBLE_LIMIT)) + else: + self.out[self.nibble_pos] |= min(extra, _NIBBLE_LIMIT) << 4 + self.nibble_pos = -1 + if extra < _NIBBLE_LIMIT: + return + extra -= _NIBBLE_LIMIT + if extra < _BYTE_LIMIT: + self.out.append(extra) + return + self.out.append(_BYTE_LIMIT) + extra += _LADDER_FLOOR + if extra < _WORD_LIMIT: + self.out += extra.to_bytes(2, "little") + else: + self.out += b"\x00\x00" + extra.to_bytes(4, "little") + + def emit_match(self, offset: int, length: int) -> None: + """Encode one match as a 16-bit LE token under a 1 flag bit.""" + token = (offset - 1) << 3 + extra = length - _MIN_MATCH + if extra < _TOKEN_LENGTH_LIMIT: + self.out += (token | extra).to_bytes(2, "little") + else: + self.out += (token | _TOKEN_LENGTH_LIMIT).to_bytes(2, "little") + self._emit_long_length(extra - _TOKEN_LENGTH_LIMIT) + self._push_flag(1) + + def finish(self) -> bytes: + """Pad the pending flag word with 1-bits and flush it.""" + pad = _FLAG_BITS - self.flag_count + self.out[self.flag_pos : self.flag_pos + 4] = (self.flags << pad | (1 << pad) - 1).to_bytes(4, "little") + return bytes(self.out) + + +def _index(table: dict[bytes, list[int]], data: bytes, pos: int) -> None: + """Record pos as a future match candidate for its 3-byte prefix.""" + key = data[pos : pos + _MIN_MATCH] + chain = table.get(key) + if chain is None: + table[key] = [pos] + return + chain.append(pos) + if len(chain) > _CHAIN_PRUNE: + del chain[:-_MAX_CHAIN] + + +def _match_length(data: bytes, candidate: int, pos: int, limit: int) -> int: + """Extend a guaranteed 3-byte prefix match as far as it goes.""" + length = _MIN_MATCH + while length < limit and data[candidate + length] == data[pos + length]: + length += 1 + return length + + +def _find_match(data: bytes, pos: int, table: dict[bytes, list[int]]) -> tuple[int, int]: + """Find the longest match for pos within the last 8192 bytes. + + Returns: + Tuple of (offset, length), or (0, 0) when no match exists. + """ + if pos + _MIN_MATCH > len(data): + return 0, 0 + chain = table.get(data[pos : pos + _MIN_MATCH]) + if chain is None: + return 0, 0 + limit = min(len(data) - pos, _MAX_MATCH_LENGTH) + best_offset = 0 + best_length = 0 + for candidate in reversed(chain[-_MAX_CHAIN:]): + if pos - candidate > _MAX_OFFSET: + break + length = _match_length(data, candidate, pos, limit) + if length > best_length: + best_offset = pos - candidate + best_length = length + if length == limit: + break + return best_offset, best_length + + +def compress(src: bytes | BinaryIO) -> bytes: + """LZXPRESS compress from a file-like object or bytes. + + Greedy hash-chain match finder per [MS-XCA] S2.3.4. + + Args: + src: File-like object or bytes to compress. + + Returns: + The compressed data. + """ + if hasattr(src, "read"): + src = src.read() + + data = bytes(src) + size = len(data) + encoder = _Encoder() + table: dict[bytes, list[int]] = {} + last_key_pos = size - _MIN_MATCH + pos = 0 + while pos < size: + offset, length = _find_match(data, pos, table) + if length: + for covered in range(pos, min(pos + length, last_key_pos + 1)): + _index(table, data, covered) + encoder.emit_match(offset, length) + pos += length + else: + if pos <= last_key_pos: + _index(table, data, pos) + encoder.emit_literal(data[pos]) + pos += 1 + return encoder.finish() diff --git a/dissect/util/compression/lzxpress9.py b/dissect/util/compression/lzxpress9.py new file mode 100644 index 0000000..97547ce --- /dev/null +++ b/dissect/util/compression/lzxpress9.py @@ -0,0 +1,1578 @@ +"""XPRESS9 compression and decompression. + +Implements the compressor and decompressor for the XPRESS9 format used by ESE +(Extensible Storage Engine) databases. XPRESS9 is an LZ77+Huffman codec with +multi-table Huffman coding, move-to-front repeated offsets, and Elias-gamma +length escaping. It has no public byte-format specification; the authority +is the MIT-licensed ESE C codec in ``dev/ese/src/_xpress9/``. + +The on-disk layout consists of one or more blocks, each starting with a 32-byte +header (magic ``0x4E86D72A``, CRC-32C protected) followed by an LSB-first +bitstream holding: optional MTF initial state, two serialized canonical-Huffman +code-length tables (short-symbol alphabet of 704, long-length alphabet of 256), +and the LZ77 token stream with Elias-gamma offset coding. + +Reference implementation: ntcompress (https://github.com/StrongWind1/ntcompress). +""" + +from __future__ import annotations + +import struct +from typing import Final + +from dissect.util.exceptions import CorruptDataError +from dissect.util.hash import crc32c as _crc32c_mod + +_crc32c = _crc32c_mod.crc32c + +# --- Block format constants (Xpress9Internal.h) --- + +BLOCK_HEADER_SIZE: Final = 32 +"""Size of ``LZ77_BLOCK_HEADER``: 8 x u32 LE (Xpress9Internal.h:972-984).""" + +_XPRESS9_MAGIC: Final = 0x4E86D72A +"""``XPRESS9_MAGIC``, word [0] of every block header (Xpress9Internal.h:970).""" + +_MAX_SHORT_LENGTH_LOG: Final = 4 +_MAX_SHORT_LENGTH: Final = 1 << _MAX_SHORT_LENGTH_LOG +_MAX_WINDOW_SIZE_LOG: Final = 24 +_MAX_MTF: Final = 4 +_LONG_LENGTH_ALPHABET_SIZE: Final = 256 +_MAX_LONG_LENGTH: Final = _LONG_LENGTH_ALPHABET_SIZE - _MAX_WINDOW_SIZE_LOG +_SHORT_SYMBOL_ALPHABET_SIZE: Final = 256 + ((_MAX_WINDOW_SIZE_LOG + _MAX_MTF) << _MAX_SHORT_LENGTH_LOG) +_MAX_CODEWORD_LENGTH: Final = 27 + +# Code-length table opcodes (Xpress9Internal.h:604-611). +_TABLE_FILL: Final = _MAX_CODEWORD_LENGTH + 1 +_TABLE_ZERO_REPT: Final = _MAX_CODEWORD_LENGTH + 2 +_TABLE_PREV: Final = _MAX_CODEWORD_LENGTH + 3 +_TABLE_ROW_0: Final = _MAX_CODEWORD_LENGTH + 4 +_TABLE_ROW_1: Final = _MAX_CODEWORD_LENGTH + 5 +_TABLE_ALPHABET_SIZE: Final = _MAX_CODEWORD_LENGTH + 6 +_TABLE_ZERO_REPT_MIN_COUNT: Final = 5 +_TABLE_FILL_BOUNDARY: Final = 16 + +_TABLE_ENCODING_STORED: Final = 0 +_TABLE_ENCODING_HUFFMAN: Final = 1 +_SMALL_TABLE_MAX_LENGTH: Final = 8 +_SMALL_TABLE_INITIAL_PREV: Final = 4 +_MAIN_TABLE_INITIAL_PREV: Final = 8 +_ZERO_REPT_EXTEND: Final = 3 +_ZERO_REPT_CONTINUE: Final = 7 + +_BLOCK_HEADER: Final = struct.Struct("<8I") + +_FLAGS_HUFFMAN_TABLE_BITS_MASK: Final = 0x1FFF +_FLAGS_RESERVED_SHIFT: Final = 20 + +_MAX_DECODED_SIZE: Final = 1 << 30 + +# --- Bit input --- + + +class _BitReader: + """LSB-first bit reader over one block's payload bitstream. + + Port of ``BIORD_*`` (Xpress9Internal.h:390-472). + """ + + __slots__ = ("_acc", "_data", "_navail", "_pos", "_size") + + _OVERRUN_BYTES: Final = 8 + + def __init__(self, data: bytes) -> None: + self._data = data + self._size = len(data) + self._pos = 0 + self._acc = 0 + self._navail = 0 + + @property + def bits_consumed(self) -> int: + """Bits consumed so far.""" + return (self._pos << 3) - self._navail + + def read(self, count: int) -> int: + """Return the next ``count`` bits as an integer, LSB first.""" + while self._navail < count: + if self._pos < self._size: + self._acc |= self._data[self._pos] << self._navail + elif self._pos >= self._size + self._OVERRUN_BYTES: + raise CorruptDataError("XPRESS9 bitstream exhausted") + self._pos += 1 + self._navail += 8 + value = self._acc & ((1 << count) - 1) + self._acc >>= count + self._navail -= count + return value + + +# --- Canonical Huffman --- + + +class _CanonicalHuffman: + """Canonical-Huffman decoder built from a code-length table. + + Equivalent to ``HuffmanCreateDecodeTables`` + ``HUFFMAN_DECODE_SYMBOL`` + (Xpress9DecHuffman.c:154-321). + """ + + __slots__ = ("_base_index", "_counts", "_first_code", "_max_length", "_single", "_symbols") + + def __init__(self, lengths: list[int]) -> None: + present = sorted((length, symbol) for symbol, length in enumerate(lengths) if length) + if not present: + raise CorruptDataError("XPRESS9 Huffman code-length table has no symbols") + + self._max_length = present[-1][0] + self._single: tuple[int, int] | None = None + if len(present) == 1: + self._single = (present[0][1], present[0][0]) + self._counts, self._first_code, self._base_index, self._symbols = [], [], [], [] + return + + counts = [0] * (self._max_length + 1) + for length, _ in present: + counts[length] += 1 + + nodes = 0 + for length in range(self._max_length, 0, -1): + nodes += counts[length] + if nodes & 1: + raise CorruptDataError(f"XPRESS9 Huffman code lengths do not form a full tree at depth {length}") + nodes >>= 1 + if nodes != 1: + raise CorruptDataError("XPRESS9 Huffman code lengths do not form a full tree") + + first_code = [0] * (self._max_length + 1) + base_index = [0] * (self._max_length + 1) + code = 0 + index = 0 + for length in range(1, self._max_length + 1): + first_code[length] = code + base_index[length] = index + code = (code + counts[length]) << 1 + index += counts[length] + + self._counts = counts + self._first_code = first_code + self._base_index = base_index + self._symbols = [symbol for _, symbol in present] + + def decode(self, reader: _BitReader) -> int: + """Decode one symbol from the stream.""" + if self._single is not None: + symbol, skip = self._single + reader.read(skip) + return symbol + + code = 0 + for length in range(1, self._max_length + 1): + code = (code << 1) | reader.read(1) + delta = code - self._first_code[length] + if delta < self._counts[length]: + return self._symbols[self._base_index[length] + delta] + + raise CorruptDataError("XPRESS9 Huffman codeword exceeds the table's maximum length") + + +# --- Code-length table decoding --- + + +def _decode_zero_run(reader: _BitReader, position: int, alphabet_size: int, fill_boundary: int) -> int: + """Read a zero-run length (Xpress9DecHuffman.c:541-564).""" + count = reader.read(2) + run = count + _TABLE_ZERO_REPT_MIN_COUNT + if count == _ZERO_REPT_EXTEND: + while True: + count = reader.read(3) + run += count + if position + run > alphabet_size or (position ^ (position + run)) >= fill_boundary: + raise CorruptDataError(f"XPRESS9 zero run of {run} at index {position} overflows") + if count != _ZERO_REPT_CONTINUE: + break + if position + run > alphabet_size or (position ^ (position + run)) >= fill_boundary: + raise CorruptDataError(f"XPRESS9 zero run of {run} at index {position} overflows") + return run + + +def _decode_small_table(reader: _BitReader) -> _CanonicalHuffman: + """Read the 33-symbol small code that encodes the length table (Xpress9DecHuffman.c:462-508).""" + prev = _SMALL_TABLE_INITIAL_PREV + lengths: list[int] = [] + for _ in range(_TABLE_ALPHABET_SIZE): + if reader.read(1) == 0: + lengths.append(prev) + continue + value = reader.read(3) + if value >= prev: + value += 1 + if value > _SMALL_TABLE_MAX_LENGTH: + raise CorruptDataError(f"XPRESS9 small-table codeword length {value} exceeds {_SMALL_TABLE_MAX_LENGTH}") + lengths.append(value) + prev = value + return _CanonicalHuffman(lengths) + + +def _apply_fill(lengths: list[int], position: int, fill_boundary: int) -> int: + """Apply a fill opcode: zeros up to the next fill boundary (Xpress9DecHuffman.c:531-539).""" + alphabet_size = len(lengths) + while True: + lengths[position] = 0 + position += 1 + if position & (fill_boundary - 1) == 0 or position >= alphabet_size: + return position + + +def _apply_row(lengths: list[int], position: int, symbol: int, fill_boundary: int) -> int: + """Apply a ROW_0/ROW_1 opcode: copy from one boundary back (Xpress9DecHuffman.c:581-613).""" + if position < fill_boundary: + raise CorruptDataError(f"XPRESS9 ROW opcode at index {position} has no previous row") + value = lengths[position - fill_boundary] + (1 if symbol == _TABLE_ROW_1 else 0) + if value == 0 or value > _MAX_CODEWORD_LENGTH: + raise CorruptDataError(f"XPRESS9 ROW opcode yields invalid codeword length {value}") + lengths[position] = value + return value + + +def _decode_coded_lengths(reader: _BitReader, alphabet_size: int, fill_boundary: int) -> list[int]: + """Read a mode-1 (Huffman-coded) code-length table (Xpress9DecHuffman.c:462-615).""" + small = _decode_small_table(reader) + lengths = [0] * alphabet_size + prev = _MAIN_TABLE_INITIAL_PREV + position = 0 + while position < alphabet_size: + symbol = small.decode(reader) + if symbol < _TABLE_FILL: + lengths[position] = symbol + position += 1 + if symbol: + prev = symbol + elif symbol == _TABLE_FILL: + position = _apply_fill(lengths, position, fill_boundary) + elif symbol == _TABLE_ZERO_REPT: + position += _decode_zero_run(reader, position, alphabet_size, fill_boundary) + elif symbol == _TABLE_PREV: + lengths[position] = prev + position += 1 + else: + prev = _apply_row(lengths, position, symbol, fill_boundary) + position += 1 + return lengths + + +def _decode_length_table(reader: _BitReader, alphabet_size: int, fill_boundary: int) -> list[int]: + """Deserialize one canonical-Huffman code-length table (Xpress9DecHuffman.c:405-622).""" + mode = reader.read(3) + if mode == _TABLE_ENCODING_STORED: + msb = alphabet_size.bit_length() - 1 + short_count = (1 << (msb + 1)) - alphabet_size + return [msb] * short_count + [msb + 1] * (alphabet_size - short_count) + if mode != _TABLE_ENCODING_HUFFMAN: + raise CorruptDataError(f"XPRESS9 unknown table encoding {mode}") + return _decode_coded_lengths(reader, alphabet_size, fill_boundary) + + +# --- Block header --- + + +def _parse_block_header( + payload: bytes | memoryview, offset: int = 0 +) -> tuple[int, int, int, int, int, int, int, int, int]: + """Parse and validate a 32-byte block header (Xpress9DecLz77.c:607-716). + + Returns a tuple of: + (orig_size, comp_size_bits, huffman_table_bits, window_size_log2, + mtf_entry_count, ptr_min_match_length, mtf_min_match_length, + session_signature, block_index) + + Raises: + CorruptDataError: Truncated, bad magic, bad CRC, or invalid flags. + """ + if len(payload) - offset < BLOCK_HEADER_SIZE: + raise CorruptDataError(f"XPRESS9 block header truncated: need {BLOCK_HEADER_SIZE} bytes") + + words = _BLOCK_HEADER.unpack_from(payload, offset) + + if words[0] != _XPRESS9_MAGIC: + raise CorruptDataError(f"Bad XPRESS9 block magic 0x{words[0]:08x}, expected 0x{_XPRESS9_MAGIC:08x}") + if words[4] != 0: + raise CorruptDataError(f"XPRESS9 block header reserved word is 0x{words[4]:08x}, must be 0") + + actual_crc = _crc32c(memoryview(payload)[offset : offset + BLOCK_HEADER_SIZE - 4]) + if words[7] != actual_crc: + raise CorruptDataError(f"XPRESS9 block header CRC-32C mismatch: 0x{words[7]:08x} vs 0x{actual_crc:08x}") + + flags = words[3] + if flags >> _FLAGS_RESERVED_SHIFT: + raise CorruptDataError(f"XPRESS9 block flags reserved bits are non-zero: 0x{flags:08x}") + + mtf_entry_count = ((flags >> 16) & 3) << 1 + if mtf_entry_count > _MAX_MTF: + raise CorruptDataError(f"XPRESS9 block declares reserved MTF entry count {mtf_entry_count}") + + huffman_table_bits = flags & _FLAGS_HUFFMAN_TABLE_BITS_MASK + if words[2] <= huffman_table_bits + BLOCK_HEADER_SIZE * 8: + raise CorruptDataError("XPRESS9 block compressed size does not exceed its header and Huffman tables") + + return ( + words[1], # orig_size + words[2], # comp_size_bits + huffman_table_bits, + ((flags >> 13) & 7) + 16, # window_size_log2 + mtf_entry_count, + ((flags >> 18) & 1) + 3, # ptr_min_match_length + ((flags >> 19) & 1) + 2, # mtf_min_match_length + words[5], # session_signature + words[6], # block_index + ) + + +# --- LZ77 token stream --- + + +def _read_mtf_initial_state(reader: _BitReader, mtf_count: int, window_log2: int) -> tuple[int, list[int]]: + """Read the MTF seed: last-was-pointer flag and initial offsets (Xpress9DecLz77.c:362-383).""" + last_was_ptr = reader.read(1) + offsets: list[int] = [] + for _ in range(mtf_count): + msb = reader.read(5) + if msb >= window_log2: + raise CorruptDataError(f"XPRESS9 MTF initial offset MSB {msb} outside {window_log2}-bit window") + offsets.append(reader.read(msb) + (1 << msb)) + return last_was_ptr, offsets + + +def _take_mtf_offset(mtf: list[int], symbol: int, *, last_was_ptr: bool) -> int: + """Pick and reorder the MTF list for one match (Xpress9Lz77Dec.i:143-220).""" + if last_was_ptr: + if symbol >= len(mtf) - 1: + raise CorruptDataError(f"XPRESS9 MTF symbol {symbol} invalid after a pointer") + offset = mtf.pop(symbol + 1) + mtf.insert(0, offset) + return offset + offset = mtf[symbol] + if symbol: + del mtf[symbol] + mtf.insert(0, offset) + return offset + + +def _read_match_length(reader: _BitReader, long_table: _CanonicalHuffman) -> int: + """Read an escaped match length from the long-length table (Xpress9Lz77Dec.i:128-141).""" + length = long_table.decode(reader) + if length >= _MAX_LONG_LENGTH: + extra_bits = length - _MAX_LONG_LENGTH + length = reader.read(extra_bits) + (1 << extra_bits) + (_MAX_LONG_LENGTH - 1) + return length + _MAX_SHORT_LENGTH - 1 + + +def _copy_match(out: bytearray, offset: int, length: int) -> None: + """Append ``length`` bytes copied from ``offset`` bytes back, with overlap semantics.""" + start = len(out) - offset + if length <= offset: + out += out[start : start + length] + else: + repeats = -(-length // offset) + out += (out[start:] * repeats)[:length] + + +def _read_block_prelude( + reader: _BitReader, + huffman_table_bits: int, + mtf_entry_count: int, + window_size_log2: int, +) -> tuple[int, list[int], _CanonicalHuffman, _CanonicalHuffman]: + """Read MTF seed and both Huffman tables before the token stream (Xpress9DecLz77.c:346-442).""" + last_was_ptr = 0 + mtf: list[int] = [] + if mtf_entry_count: + last_was_ptr, mtf = _read_mtf_initial_state(reader, mtf_entry_count, window_size_log2) + short_table = _CanonicalHuffman(_decode_length_table(reader, _SHORT_SYMBOL_ALPHABET_SIZE, _MAX_SHORT_LENGTH)) + long_table = _CanonicalHuffman(_decode_length_table(reader, _LONG_LENGTH_ALPHABET_SIZE, _MAX_SHORT_LENGTH)) + if reader.bits_consumed != huffman_table_bits: + raise CorruptDataError( + f"XPRESS9 Huffman tables consumed {reader.bits_consumed} bits, header declares {huffman_table_bits}" + ) + return last_was_ptr, mtf, short_table, long_table + + +def _decode_block( + reader: _BitReader, + orig_size: int, + comp_size_bits: int, + huffman_table_bits: int, + mtf_entry_count: int, + ptr_min_match_length: int, + mtf_min_match_length: int, + window_size_log2: int, + out: bytearray, +) -> None: + """Decode one block's bitstream into ``out`` (Xpress9Lz77Dec.i:101-293).""" + last_was_ptr, mtf, short_table, long_table = _read_block_prelude( + reader, huffman_table_bits, mtf_entry_count, window_size_log2 + ) + target = len(out) + orig_size + + while len(out) < target: + symbol = short_table.decode(reader) + if symbol < 256: + out.append(symbol) + last_was_ptr = 0 + continue + + symbol -= 256 + length = symbol & (_MAX_SHORT_LENGTH - 1) + symbol >>= _MAX_SHORT_LENGTH_LOG + + if length == _MAX_SHORT_LENGTH - 1: + length = _read_match_length(reader, long_table) + + if symbol < mtf_entry_count: + length += mtf_min_match_length + offset = _take_mtf_offset(mtf, symbol, last_was_ptr=bool(last_was_ptr)) + else: + length += ptr_min_match_length + msb = symbol - mtf_entry_count + offset = reader.read(msb) + (1 << msb) + if mtf: + mtf.insert(0, offset) + mtf.pop() + + if offset > len(out): + raise CorruptDataError(f"XPRESS9 match offset {offset} reaches before the start of data") + if length > target - len(out): + raise CorruptDataError(f"XPRESS9 match of length {length} overruns the block's declared size") + + last_was_ptr = 1 + _copy_match(out, offset, length) + + if len(out) != target: + raise CorruptDataError( + f"XPRESS9 block decoded {len(out) - target + orig_size} bytes, header declares {orig_size}" + ) + actual_bits = reader.bits_consumed + BLOCK_HEADER_SIZE * 8 + if actual_bits != comp_size_bits: + raise CorruptDataError(f"XPRESS9 block consumed {actual_bits} bits, header declares {comp_size_bits}") + + +# --- Session (block sequence) --- + + +def _session_blocks(payload: bytes) -> list[tuple[tuple[int, ...], int, int]]: + """Split a codec payload into validated, byte-aligned blocks. + + Returns ``(header_fields, body_start, body_end)`` per block. + + Raises: + CorruptDataError: Empty payload, truncated block, or header validation failure. + """ + blocks: list[tuple[tuple[int, ...], int, int]] = [] + first: tuple[int, ...] | None = None + offset = 0 + while offset < len(payload): + header = _parse_block_header(payload, offset) + if first is None: + first = header + elif header[7] != first[7]: + raise CorruptDataError("XPRESS9 session signature mismatch between blocks") + elif header[3:7] != first[3:7]: + raise CorruptDataError("XPRESS9 block changes the session's coding parameters") + if header[8] != len(blocks): + raise CorruptDataError(f"XPRESS9 block index {header[8]}, expected {len(blocks)}") + + end = offset + ((header[1] + 7) >> 3) + if end > len(payload): + raise CorruptDataError(f"XPRESS9 block truncated: needs {end - offset} bytes") + blocks.append((header, offset + BLOCK_HEADER_SIZE, end)) + offset = end + + if not blocks: + raise CorruptDataError("XPRESS9 payload has no blocks") + return blocks + + +def decompress(src: bytes) -> bytes: + """Decompress raw XPRESS9 block data. + + Takes the codec payload (the bytes *after* any ESE record header) and + returns the decompressed plaintext. The payload consists of one or more + self-describing XPRESS9 blocks, each with a 32-byte header and a bitstream. + + Args: + src: Raw XPRESS9 block data. + + Returns: + The decompressed data. + + Raises: + CorruptDataError: If the data is truncated, corrupt, or violates format invariants. + """ + blocks = _session_blocks(src) + declared_total = sum(h[0] for h, _, _ in blocks) + if declared_total > _MAX_DECODED_SIZE: + raise CorruptDataError(f"XPRESS9 declares {declared_total} bytes, over the safety limit") + + out = bytearray() + for header, body_start, body_end in blocks: + _decode_block( + _BitReader(src[body_start:body_end]), + header[0], # orig_size + header[1], # comp_size_bits + header[2], # huffman_table_bits + header[4], # mtf_entry_count + header[5], # ptr_min_match_length + header[6], # mtf_min_match_length + header[3], # window_size_log2 + out, + ) + return bytes(out) + + +def decompressed_size(src: bytes) -> int: + """Return the total decompressed size from block headers, without decoding. + + Args: + src: Raw XPRESS9 block data. + + Returns: + Sum of all block ``orig_size`` fields. + + Raises: + CorruptDataError: If the block headers are truncated or corrupt. + """ + return sum(h[0] for h, _, _ in _session_blocks(src)) + + +# --- Encoder constants (Xpress9EncLz77.c) --- + +# ESE fixed "Cosmos Level 6" parameters (compression.cxx:1696-1706). +_ENC_WINDOW_SIZE_LOG2: Final = 16 +_ENC_WINDOW_SIZE: Final = 1 << _ENC_WINDOW_SIZE_LOG2 +_ENC_MTF_ENTRY_COUNT: Final = 4 +_ENC_PTR_MIN_MATCH_LENGTH: Final = 4 +_ENC_MTF_MIN_MATCH_LENGTH: Final = 2 +_ENC_SESSION_SIGNATURE: Final = 0x12345678 + +# Hash table sizing (Xpress9EncLz77.c:1062-1102). +_HASH_TABLE_SIZE_LOG2: Final = 12 +_HASH_TABLE_SIZE: Final = 1 << _HASH_TABLE_SIZE_LOG2 +_HASH_MASK: Final = _HASH_TABLE_SIZE - 1 + +# Lookup depth: LookupDepth=1, +1 before the loop (Xpress9Lz77EncPass1.i:65). +_ENC_LOOKUP_DEPTH: Final = 1 +_ENC_MAX_DEPTH: Final = _ENC_LOOKUP_DEPTH + 1 + +# IR buffer chunk size (Xpress9EncLz77.c:926,1570-1581). +_IR_BUFFER_SIZE: Final = 2 * _ENC_WINDOW_SIZE +_IR_CHUNK_SIZE: Final = _IR_BUFFER_SIZE // 3 - 256 + +# Maximum-offset-by-length table (Xpress9EncLz77.c:297-310). +_MAX_OFFSET_BY_LENGTH: Final[tuple[int, ...]] = ( + 0, + 0, + -(1 << 6), + -(1 << 10), + -(1 << 13), + -(1 << 16), +) + +# Token types for the intermediate representation. +_TOKEN_LIT: Final = 0 +_TOKEN_PTR: Final = 1 +_TOKEN_MTF: Final = 2 + + +# --- Bit output (encoder) --- + + +class _BitWriter: + """LSB-first bit writer, the encoding counterpart of ``_BitReader``. + + Port of the ``BIOWR`` macro family (Xpress9Internal.h:313-372). + """ + + __slots__ = ("_acc", "_buf", "_navail") + + def __init__(self) -> None: + self._buf = bytearray() + self._acc = 0 + self._navail = 0 + + @property + def bits_written(self) -> int: + """Total bits written so far, including unflushed accumulator bits.""" + return len(self._buf) * 8 + self._navail + + def write(self, value: int, count: int) -> None: + """Write ``count`` bits of ``value``, LSB first.""" + self._acc |= value << self._navail + self._navail += count + while self._navail >= 8: + self._buf.append(self._acc & 0xFF) + self._acc >>= 8 + self._navail -= 8 + + def flush(self) -> None: + """Write any remaining partial byte.""" + if self._navail > 0: + self._buf.append(self._acc & 0xFF) + self._acc = 0 + self._navail = 0 + + def getvalue(self) -> bytes: + """Return the written bytes. Must call ``flush`` first.""" + return bytes(self._buf) + + +# --- Canonical Huffman encoder --- + + +def _reverse_mask(value: int, bits: int) -> int: + """Reverse the lowest ``bits`` bits (HuffmanReverseMask, Xpress9Misc.c:123-149).""" + result = 0 + for _ in range(bits): + result = (result << 1) | (value & 1) + value >>= 1 + return result + + +def _compute_depths_from_tree( + queue_leaf: list[tuple[int, int]], + _queue_node: list[tuple[int, list[int]]], + depths: list[int], + _li: int, + _ni: int, + n: int, +) -> None: + """Compute symbol depths by building and walking the Huffman tree.""" + if n <= 1: + if n == 1: + depths[queue_leaf[0][1]] = 1 + return + + sorted_symbols = [s for _, s in sorted(queue_leaf, key=lambda x: (x[0], x[1]))] + sorted_counts = [c for c, _ in sorted(queue_leaf, key=lambda x: (x[0], x[1]))] + + class _Node: + __slots__ = ("count", "depth", "left", "right", "symbol") + + def __init__(self, *, count: int, symbol: int = -1) -> None: + self.count = count + self.symbol = symbol + self.depth = 0 + self.left: _Node | None = None + self.right: _Node | None = None + + leaves = [_Node(count=sorted_counts[i], symbol=sorted_symbols[i]) for i in range(n)] + internals: list[_Node] = [] + ssi = 0 + isi = 0 + + def _pop_min_node() -> _Node: + nonlocal ssi, isi + sv = leaves[ssi].count if ssi < n else 2**63 + iv = internals[isi].count if isi < len(internals) else 2**63 + if sv <= iv: + nd = leaves[ssi] + ssi += 1 + return nd + nd = internals[isi] + isi += 1 + return nd + + for _ in range(n - 1): + left = _pop_min_node() + right = _pop_min_node() + parent = _Node(count=left.count + right.count) + parent.left = left + parent.right = right + internals.append(parent) + + root = internals[-1] if internals else leaves[0] + + stack: list[tuple[_Node, int]] = [(root, 0)] + while stack: + node, d = stack.pop() + if node.symbol >= 0: + depths[node.symbol] = max(d, 1) + else: + if node.left is not None: + stack.append((node.left, d + 1)) + if node.right is not None: + stack.append((node.right, d + 1)) + + +def _build_huffman_codes(counts: list[int], alphabet_size: int, max_codeword_length: int) -> list[tuple[int, int]]: + """Build canonical Huffman codewords from frequency counts. + + Port of ``Xpress9HuffmanCreateTree`` (Xpress9EncHuffman.c:673-797): build a tree, + compute bit-lengths, truncate to ``max_codeword_length``, canonize, and create + bit-reversed codewords. Returns ``(reversed_codeword, length)`` per symbol. + """ + present = sorted( + ((counts[i], i) for i in range(alphabet_size) if counts[i] > 0), + key=lambda x: (x[0], x[1]), + ) + n = len(present) + result: list[tuple[int, int]] = [(0, 0)] * alphabet_size + + if n == 0: + return result + + if n == 1: + sym = present[0][1] + result[sym] = (0, 1) + return result + + queue_leaf = list(present) + queue_node: list[tuple[int, list[int]]] = [] + li = 0 + ni = 0 + + def _pop_min() -> tuple[int, list[int]]: + nonlocal li, ni + lv = queue_leaf[li][0] if li < len(queue_leaf) else 2**63 + nv = queue_node[ni][0] if ni < len(queue_node) else 2**63 + if lv <= nv: + c, s = queue_leaf[li] + li += 1 + return (c, [s]) + r = queue_node[ni] + ni += 1 + return r + + for _ in range(n - 1): + left_c, left_s = _pop_min() + right_c, right_s = _pop_min() + queue_node.append((left_c + right_c, left_s + right_s)) + + depths = [0] * alphabet_size + _compute_depths_from_tree(queue_leaf, queue_node, depths, li, ni, n) + + bit_length_count = [0] * 65 + for d in depths: + if d > 0: + bit_length_count[d] += 1 + + # Truncate tree to max_codeword_length (HuffmanTruncateTree, Xpress9EncHuffman.c:460-508). + for i in range(64, max_codeword_length, -1): + while bit_length_count[i] > 0: + j = max_codeword_length - 1 + while j > 0 and bit_length_count[j] == 0: + j -= 1 + bit_length_count[j] -= 1 + bit_length_count[j + 1] += 2 + bit_length_count[i - 1] += 1 + bit_length_count[i] -= 2 + + # Assign final lengths (HuffmanCanonizeTree, Xpress9EncHuffman.c:514-617). + sorted_symbols = [s for _, s in present] + sym_lengths = [0] * alphabet_size + idx = 0 + for length in range(max_codeword_length, 0, -1): + for _ in range(bit_length_count[length]): + sym_lengths[sorted_symbols[idx]] = length + idx += 1 + + # Build canonical codewords (HuffmanCreateCodewords, Xpress9EncHuffman.c:624-665). + by_length_symbol = sorted( + ((sym_lengths[s], s) for s in range(alphabet_size) if sym_lengths[s] > 0), + key=lambda x: (x[0], x[1]), + ) + code = 0 + prev_len = 0 + for length, sym in by_length_symbol: + code <<= length - prev_len + result[sym] = (_reverse_mask(code, length), length) + code += 1 + prev_len = length + + return result + + +# --- Huffman table serialization (encoder) --- + + +def _encode_small_table(writer: _BitWriter, code_lengths: list[tuple[int, int]]) -> None: + """Write the 33-symbol small code-length table (Xpress9EncHuffman.c:941-964).""" + prev = _SMALL_TABLE_INITIAL_PREV + for sym in range(_TABLE_ALPHABET_SIZE): + length = code_lengths[sym][1] + if length == prev: + writer.write(0, 1) + else: + writer.write(1, 1) + if length > prev: + writer.write(length - 1, 3) + else: + writer.write(length, 3) + prev = length + + +def _write_length_sequence( + writer: _BitWriter, + lengths: list[int], + symbols: list[int], + meta_codes: list[tuple[int, int]], + alphabet_size: int, + fill_boundary: int, +) -> None: + """Write the Mode 1 encoded length sequence into ``writer``.""" + prev_sym = _MAIN_TABLE_INITIAL_PREV + sym_idx = 0 + i = 0 + while i < alphabet_size: + k = lengths[i] + if k != 0: + meta_sym = symbols[sym_idx] + sym_idx += 1 + cw, cl = meta_codes[meta_sym] + writer.write(cw, cl) + if k != prev_sym: + prev_sym = k + i += 1 + else: + zero_start = i + while i < alphabet_size and lengths[i] == 0: + i += 1 + k_pos = zero_start + while (k_pos ^ i) >= fill_boundary: + cw, cl = meta_codes[_TABLE_FILL] + writer.write(cw, cl) + sym_idx += 1 + k_pos = (k_pos & ~(fill_boundary - 1)) + fill_boundary + remaining = i - k_pos + if remaining > 0: + if remaining < _TABLE_ZERO_REPT_MIN_COUNT: + for _ in range(remaining): + cw, cl = meta_codes[0] + writer.write(cw, cl) + sym_idx += 1 + else: + cw, cl = meta_codes[_TABLE_ZERO_REPT] + writer.write(cw, cl) + sym_idx += 1 + run = remaining - _TABLE_ZERO_REPT_MIN_COUNT + if run < 3: + writer.write(run, 2) + else: + writer.write(3, 2) + run -= 3 + while run >= 7: + writer.write(7, 3) + run -= 7 + writer.write(run, 3) + + +def _encode_huffman_table( + writer: _BitWriter, + codes: list[tuple[int, int]], + counts: list[int], + alphabet_size: int, + fill_boundary: int, +) -> list[tuple[int, int]]: + """Serialize one canonical-Huffman code-length table into the bitstream. + + Port of ``Xpress9HuffmanCreateAndEncodeTable`` (Xpress9EncHuffman.c:1078-1222): + compare mode 0 (stored/uniform) vs mode 1 (Huffman-coded) and pick the cheaper one. + """ + lengths = [c[1] for c in codes] + + # Mode 0 (stored/uniform) cost. + msb = alphabet_size.bit_length() - 1 + threshold = (1 << (msb + 1)) - alphabet_size + freq0 = sum(counts[i] for i in range(threshold)) + freq1 = sum(counts[i] for i in range(threshold, alphabet_size)) + mode0_cost = (freq0 + freq1) * msb + freq1 + 3 + + # Mode 1 (Huffman-coded) data cost. + huffman_data_cost = sum(counts[i] * lengths[i] for i in range(alphabet_size) if lengths[i] > 0) + + # Build Mode 1 opcode sequence. + symbols: list[int] = [] + meta_counts = [0] * _TABLE_ALPHABET_SIZE + prev_sym = _MAIN_TABLE_INITIAL_PREV + i = 0 + while i < alphabet_size: + k = lengths[i] + if k != 0: + sym: int + if k == prev_sym: + sym = _TABLE_PREV + else: + prev_sym = k + if i >= fill_boundary: + row_val = lengths[i - fill_boundary] + if k == row_val: + sym = _TABLE_ROW_0 + elif k == row_val + 1: + sym = _TABLE_ROW_1 + else: + sym = k + else: + sym = k + meta_counts[sym] += 1 + symbols.append(sym) + i += 1 + else: + zero_start = i + while i < alphabet_size and lengths[i] == 0: + i += 1 + k_pos = zero_start + while (k_pos ^ i) >= fill_boundary: + symbols.append(_TABLE_FILL) + meta_counts[_TABLE_FILL] += 1 + k_pos = (k_pos & ~(fill_boundary - 1)) + fill_boundary + remaining = i - k_pos + if remaining > 0: + if remaining < _TABLE_ZERO_REPT_MIN_COUNT: + for _ in range(remaining): + symbols.append(0) + meta_counts[0] += 1 + else: + symbols.append(_TABLE_ZERO_REPT) + meta_counts[_TABLE_ZERO_REPT] += 1 + + meta_codes = _build_huffman_codes(meta_counts, _TABLE_ALPHABET_SIZE, _SMALL_TABLE_MAX_LENGTH) + + # Measure Mode 1 cost exactly by writing to a scratch writer + # (Xpress9EncHuffman.c:1129-1141). + scratch = _BitWriter() + scratch.write(_TABLE_ENCODING_HUFFMAN, 3) + _encode_small_table(scratch, meta_codes) + _write_length_sequence(scratch, lengths, symbols, meta_codes, alphabet_size, fill_boundary) + mode1_cost = huffman_data_cost + scratch.bits_written + + if mode0_cost <= mode1_cost: + writer.write(_TABLE_ENCODING_STORED, 3) + uniform: list[tuple[int, int]] = [(0, 0)] * alphabet_size + for j in range(threshold): + uniform[j] = (_reverse_mask(j, msb), msb) + base = threshold << 1 + for j in range(threshold, alphabet_size): + uniform[j] = (_reverse_mask(base, msb + 1), msb + 1) + base += 1 + return uniform + + writer.write(_TABLE_ENCODING_HUFFMAN, 3) + _encode_small_table(writer, meta_codes) + _write_length_sequence(writer, lengths, symbols, meta_codes, alphabet_size, fill_boundary) + return codes + + +# --- LZ77 match finder and encoder --- + + +def _hash4(data: bytes, pos: int) -> int: + """Hash 4 bytes at ``pos`` (Xpress9Lz77EncInsert.i:212-225, non-SSE2 scalar path).""" + if pos + 4 > len(data): + return 0 + v = data[pos] | (data[pos + 1] << 8) | (data[pos + 2] << 16) | (data[pos + 3] << 24) + v = ((v ^ 0xDEADBEEF) + (v >> 5)) & 0xFFFFFFFF + v = (v ^ (v >> 11)) & 0xFFFFFFFF + return v & _HASH_MASK + + +def _hash_insert(data: bytes, pos: int, hash_table: list[int], p_next: list[int]) -> None: + """Insert ``pos`` into the hash chain (Xpress9Lz77EncInsert.i:226-229).""" + h = _hash4(data, pos) + p_next[pos] = hash_table[h] + hash_table[h] = pos + + +def _hash_insert_range( + data: bytes, start: int, end: int, hash_table: list[int], p_next: list[int], data_size: int +) -> None: + """Insert all positions in ``[start, end)`` into the hash chain.""" + hashable_limit = data_size - 4 if data_size >= 4 else 0 + for pos in range(start, end): + if pos < hashable_limit: + _hash_insert(data, pos, hash_table, p_next) + else: + p_next[pos] = 0 + + +def _chain_lookup( + data: bytes, pos: int, data_size: int, p_next: list[int], max_depth: int, best_len: int +) -> tuple[int, int]: + """Walk the hash chain from ``pos`` looking for the longest match. + + Port of the Xpress9Lookup.i inner loop (DEEP_LOOKUP=1, TAIL_T=UInt16). + Returns ``(best_offset_neg, best_length)``. + """ + candidate = p_next[pos] + saved_next_0 = p_next[0] + p_next[0] = pos + + best_offset = 0 + depth_remaining = max_depth + max_offset_table = _MAX_OFFSET_BY_LENGTH + max_offset_len = len(max_offset_table) - 1 + + while True: + tail_pos = pos + best_len - 1 + if tail_pos + 1 >= data_size: + break + tail0 = data[tail_pos] + tail1 = data[tail_pos + 1] + + found_candidate = -1 + checks_in_block = 0 + cur = candidate + while checks_in_block < 8: + next_cur = p_next[cur] + if cur + best_len < data_size and data[cur + best_len - 1] == tail0 and data[cur + best_len] == tail1: + found_candidate = cur + candidate = next_cur + break + checks_in_block += 1 + + cur2 = p_next[next_cur] + if ( + next_cur + best_len < data_size + and data[next_cur + best_len - 1] == tail0 + and data[next_cur + best_len] == tail1 + ): + found_candidate = next_cur + candidate = cur2 + break + checks_in_block += 1 + cur = cur2 + + if found_candidate < 0: + depth_remaining -= 1 + if depth_remaining == 0: + break + candidate = cur + continue + + if found_candidate >= pos: + break + + i_offset = found_candidate - pos + ml = 0 + while pos + ml < data_size and data[pos + ml] == data[found_candidate + ml]: + ml += 1 + + if ml > best_len and (ml > max_offset_len or i_offset > max_offset_table[ml]): + best_len = ml + best_offset = i_offset + + candidate = p_next[found_candidate] + + p_next[0] = saved_next_0 + return best_offset, best_len + + +def _check_mtf(data: bytes, pos: int, i_offset: int, data_size: int) -> int: + """Check an MTF match at ``pos`` and return its length (Xpress9EncLz77.c:42-66).""" + if pos + i_offset <= 0: + return 0 + src = pos + i_offset + if data[pos] != data[src] or data[pos + 1] != data[src + 1]: + return 0 + ml = _ENC_MTF_MIN_MATCH_LENGTH + while pos + ml < data_size and data[pos + ml] == data[src + ml]: + ml += 1 + return ml + + +def _emit_mtf( + tokens: list[tuple[int, ...]], + encode_idx: int, + offset_neg: int, + length: int, + mtf_slot: int, + mtf_0: int, + mtf_1: int, + mtf_2: int, + mtf_3: int, +) -> tuple[int, int, int, int]: + """Emit an MTF token and perform UPDATE_MTF, returning the updated MTF state.""" + tokens.append((_TOKEN_MTF, encode_idx, -offset_neg, length)) + if mtf_slot >= 3: + mtf_3 = mtf_2 + if mtf_slot >= 2: + mtf_2 = mtf_1 + mtf_1 = mtf_0 + mtf_0 = offset_neg + return mtf_0, mtf_1, mtf_2, mtf_3 + + +def _check_all_mtf( + data: bytes, + pos: int, + data_size: int, + mtf_last_ptr: int, + mtf_0: int, + mtf_1: int, + mtf_2: int, + mtf_3: int, +) -> tuple[int, int, int]: + """Check all 4 MTF slots at ``pos`` (Xpress9Lz77EncPass1.i:80-93).""" + if mtf_last_ptr == 0: + ml = _check_mtf(data, pos, mtf_0, data_size) + if ml >= _ENC_MTF_MIN_MATCH_LENGTH: + return ml, 0, 0 + ml = _check_mtf(data, pos, mtf_1, data_size) + if ml >= _ENC_MTF_MIN_MATCH_LENGTH: + return ml, 1 + mtf_last_ptr, 1 + ml = _check_mtf(data, pos, mtf_2, data_size) + if ml >= _ENC_MTF_MIN_MATCH_LENGTH: + return ml, 2 + mtf_last_ptr, 2 + ml = _check_mtf(data, pos, mtf_3, data_size) + if ml >= _ENC_MTF_MIN_MATCH_LENGTH: + return ml, 3 + mtf_last_ptr, 3 + return 0, 0, 0 + + +def _lookahead_check_mtf( + data: bytes, + pos: int, + data_size: int, + threshold: int, + mtf_0: int, + mtf_1: int, + mtf_2: int, + mtf_3: int, +) -> tuple[int, int, int]: + """Check all 4 MTF slots for lookahead paths (Xpress9EncLz77.c:71-136).""" + for slot, off in enumerate((mtf_0, mtf_1, mtf_2, mtf_3)): + ml = _check_mtf(data, pos, off, data_size) + if ml >= threshold: + return ml, slot, slot + return 0, 0, 0 + + +def _lz77_tokenize(data: bytes) -> list[tuple[int, ...]]: + """LZ77 tokenizer with lazy match evaluation matching the C reference encoder. + + Port of Xpress9Lz77EncPass1.i (DEEP_LOOKUP=1, LZ77_MTF=4, LAZY_MATCH_EVALUATION). + """ + data_size = len(data) + tokens: list[tuple[int, ...]] = [] + if data_size == 0: + return tokens + + hash_table = [0] * _HASH_TABLE_SIZE + p_next = [0] * data_size + + hash_insert_pos = 0 + mtf_0 = -1 + mtf_1 = -1 + mtf_2 = -1 + mtf_3 = -1 + mtf_last_ptr = 0 + + max_depth = _ENC_MAX_DEPTH + pos = 0 + bytes_copied = 0 + + while bytes_copied < data_size: + remaining = data_size - bytes_copied + chunk = min(_IR_CHUNK_SIZE, remaining) + bytes_copied += chunk + chunk_data_size = bytes_copied + + _hash_insert_range(data, hash_insert_pos, bytes_copied, hash_table, p_next, data_size) + hash_insert_pos = min(bytes_copied, max(data_size - 4, 0)) if data_size >= 4 else 0 + stop_position = hash_insert_pos + + while pos < stop_position: + # MTF checks. + mtf_ml, mtf_ei, mtf_sl = _check_all_mtf( + data, pos, chunk_data_size, mtf_last_ptr, mtf_0, mtf_1, mtf_2, mtf_3 + ) + if mtf_ml > 0: + offset_neg = (mtf_0, mtf_1, mtf_2, mtf_3)[mtf_sl] + mtf_0, mtf_1, mtf_2, mtf_3 = _emit_mtf( + tokens, mtf_ei, offset_neg, mtf_ml, mtf_sl, mtf_0, mtf_1, mtf_2, mtf_3 + ) + mtf_last_ptr = -1 + pos += mtf_ml + continue + + # Hash chain lookup. + if p_next[pos] != 0: + best_len = _ENC_PTR_MIN_MATCH_LENGTH - 1 + best_offset, best_len = _chain_lookup(data, pos, chunk_data_size, p_next, max_depth, best_len) + + if best_len >= _ENC_PTR_MIN_MATCH_LENGTH: + # LAZY_MATCH_EVALUATION. + if pos + 2 < stop_position and p_next[pos + 1] != 0: + saved_best_len = best_len + saved_best_offset = best_offset + + la_threshold = max(saved_best_len - 3, _ENC_MTF_MIN_MATCH_LENGTH) + la_ml, la_ei, la_sl = _lookahead_check_mtf( + data, + pos + 1, + chunk_data_size, + la_threshold, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + if la_ml >= la_threshold: + tokens.append((_TOKEN_LIT, data[pos])) + offset_neg = (mtf_0, mtf_1, mtf_2, mtf_3)[la_sl] + mtf_0, mtf_1, mtf_2, mtf_3 = _emit_mtf( + tokens, + la_ei, + offset_neg, + la_ml, + la_sl, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + mtf_last_ptr = -1 + pos += 1 + la_ml + continue + + la1_best_len = best_len + la1_best_offset, la1_best_len = _chain_lookup( + data, + pos + 1, + chunk_data_size, + p_next, + max(max_depth // 2, 1), + la1_best_len, + ) + + if la1_best_len > saved_best_len: + if p_next[pos + 2] != 0: + saved_best_len2 = la1_best_len + saved_best_offset2 = la1_best_offset + + la2_threshold = max(saved_best_len2 - 3, _ENC_MTF_MIN_MATCH_LENGTH) + la2_ml, la2_ei, la2_sl = _lookahead_check_mtf( + data, + pos + 2, + chunk_data_size, + la2_threshold, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + if la2_ml >= la2_threshold: + tokens.append((_TOKEN_LIT, data[pos])) + tokens.append((_TOKEN_LIT, data[pos + 1])) + offset_neg = (mtf_0, mtf_1, mtf_2, mtf_3)[la2_sl] + mtf_0, mtf_1, mtf_2, mtf_3 = _emit_mtf( + tokens, + la2_ei, + offset_neg, + la2_ml, + la2_sl, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + mtf_last_ptr = -1 + pos += 2 + la2_ml + continue + + la2_best_len = la1_best_len + la2_best_offset, la2_best_len = _chain_lookup( + data, + pos + 2, + chunk_data_size, + p_next, + max(max_depth // 4, 1), + la2_best_len, + ) + + if la2_best_len > saved_best_len2: + tokens.append((_TOKEN_LIT, data[pos])) + tokens.append((_TOKEN_LIT, data[pos + 1])) + pos += 2 + best_offset = la2_best_offset + best_len = la2_best_len + else: + tokens.append((_TOKEN_LIT, data[pos])) + pos += 1 + best_offset = saved_best_offset2 + best_len = saved_best_len2 + else: + tokens.append((_TOKEN_LIT, data[pos])) + pos += 1 + best_offset = la1_best_offset + best_len = la1_best_len + else: + la2_threshold = max(saved_best_len - 3, _ENC_MTF_MIN_MATCH_LENGTH) + la2_ml, la2_ei, la2_sl = _lookahead_check_mtf( + data, + pos + 2, + chunk_data_size, + la2_threshold, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + if la2_ml >= la2_threshold: + tokens.append((_TOKEN_LIT, data[pos])) + tokens.append((_TOKEN_LIT, data[pos + 1])) + offset_neg = (mtf_0, mtf_1, mtf_2, mtf_3)[la2_sl] + mtf_0, mtf_1, mtf_2, mtf_3 = _emit_mtf( + tokens, + la2_ei, + offset_neg, + la2_ml, + la2_sl, + mtf_0, + mtf_1, + mtf_2, + mtf_3, + ) + mtf_last_ptr = -1 + pos += 2 + la2_ml + continue + + best_offset = saved_best_offset + best_len = saved_best_len + + # Emit pointer match. + tokens.append((_TOKEN_PTR, -best_offset, best_len)) + mtf_3 = mtf_2 + mtf_2 = mtf_1 + mtf_1 = mtf_0 + mtf_0 = best_offset + mtf_last_ptr = -1 + pos += best_len + continue + + # Literal loop (Xpress9Lz77EncPass1.i:243-267). + mtf_last_ptr = 0 + while True: + tokens.append((_TOKEN_LIT, data[pos])) + pos += 1 + if pos >= stop_position: + break + cand = p_next[pos] + if cand == 0: + continue + if ( + pos + 3 < chunk_data_size + and cand + 3 < chunk_data_size + and data[pos] == data[cand] + and data[pos + 1] == data[cand + 1] + and data[pos + 2] == data[cand + 2] + and data[pos + 3] == data[cand + 3] + ): + break + + # Flush trailing unhashed positions as literals. + while pos < data_size: + tokens.append((_TOKEN_LIT, data[pos])) + mtf_last_ptr = 0 + pos += 1 + + return tokens + + +# --- Token encoding --- + + +def _collect_frequencies(tokens: list[tuple[int, ...]]) -> tuple[list[int], list[int], int]: + """Collect Huffman frequency tables from the token stream.""" + short_counts = [0] * _SHORT_SYMBOL_ALPHABET_SIZE + long_counts = [0] * _LONG_LENGTH_ALPHABET_SIZE + extra_bits = 0 + + for token in tokens: + if token[0] == _TOKEN_LIT: + short_counts[token[1]] += 1 + elif token[0] == _TOKEN_PTR: + offset, length = token[1], token[2] + msb_offset = offset.bit_length() - 1 + extra_bits += msb_offset + sym_base = (msb_offset + 16 + _ENC_MTF_ENTRY_COUNT) << _MAX_SHORT_LENGTH_LOG + adj_length = length - _ENC_PTR_MIN_MATCH_LENGTH + if adj_length < _MAX_SHORT_LENGTH - 1: + short_counts[sym_base + adj_length] += 1 + else: + short_counts[sym_base + _MAX_SHORT_LENGTH - 1] += 1 + long_length = adj_length - (_MAX_SHORT_LENGTH - 1) + if long_length <= _MAX_LONG_LENGTH - 1: + long_counts[long_length] += 1 + else: + escaped = long_length - (_MAX_LONG_LENGTH - 1) + msb_len = escaped.bit_length() - 1 + extra_bits += msb_len + long_counts[msb_len + _MAX_LONG_LENGTH] += 1 + else: # _TOKEN_MTF + mtf_index, _offset, length = token[1], token[2], token[3] + sym_base = (mtf_index + 16) << _MAX_SHORT_LENGTH_LOG + adj_length = length - _ENC_MTF_MIN_MATCH_LENGTH + if adj_length < _MAX_SHORT_LENGTH - 1: + short_counts[sym_base + adj_length] += 1 + else: + short_counts[sym_base + _MAX_SHORT_LENGTH - 1] += 1 + long_length = adj_length - (_MAX_SHORT_LENGTH - 1) + if long_length <= _MAX_LONG_LENGTH - 1: + long_counts[long_length] += 1 + else: + escaped = long_length - (_MAX_LONG_LENGTH - 1) + msb_len = escaped.bit_length() - 1 + extra_bits += msb_len + long_counts[msb_len + _MAX_LONG_LENGTH] += 1 + + return short_counts, long_counts, extra_bits + + +def _encode_tokens( + writer: _BitWriter, + tokens: list[tuple[int, ...]], + short_codes: list[tuple[int, int]], + long_codes: list[tuple[int, int]], +) -> None: + """Encode the LZ77 token stream using Huffman codes (Xpress9Lz77EncPass2.i:1-107).""" + for token in tokens: + if token[0] == _TOKEN_LIT: + cw, cl = short_codes[token[1]] + writer.write(cw, cl) + elif token[0] == _TOKEN_PTR: + offset, length = token[1], token[2] + msb_offset = offset.bit_length() - 1 + low_offset = offset - (1 << msb_offset) + sym_base = (msb_offset + 16 + _ENC_MTF_ENTRY_COUNT) << _MAX_SHORT_LENGTH_LOG + adj_length = length - _ENC_PTR_MIN_MATCH_LENGTH + + if adj_length < _MAX_SHORT_LENGTH - 1: + cw, cl = short_codes[sym_base + adj_length] + writer.write(cw, cl) + else: + cw, cl = short_codes[sym_base + _MAX_SHORT_LENGTH - 1] + writer.write(cw, cl) + long_length = adj_length - (_MAX_SHORT_LENGTH - 1) + if long_length <= _MAX_LONG_LENGTH - 1: + cw, cl = long_codes[long_length] + writer.write(cw, cl) + else: + escaped = long_length - (_MAX_LONG_LENGTH - 1) + msb_len = escaped.bit_length() - 1 + low_len = escaped - (1 << msb_len) + cw, cl = long_codes[msb_len + _MAX_LONG_LENGTH] + writer.write(cw, cl) + writer.write(low_len, msb_len) + + writer.write(low_offset, msb_offset) + + else: # _TOKEN_MTF + mtf_index, _offset, length = token[1], token[2], token[3] + sym_base = (mtf_index + 16) << _MAX_SHORT_LENGTH_LOG + adj_length = length - _ENC_MTF_MIN_MATCH_LENGTH + + if adj_length < _MAX_SHORT_LENGTH - 1: + cw, cl = short_codes[sym_base + adj_length] + writer.write(cw, cl) + else: + cw, cl = short_codes[sym_base + _MAX_SHORT_LENGTH - 1] + writer.write(cw, cl) + long_length = adj_length - (_MAX_SHORT_LENGTH - 1) + if long_length <= _MAX_LONG_LENGTH - 1: + cw, cl = long_codes[long_length] + writer.write(cw, cl) + else: + escaped = long_length - (_MAX_LONG_LENGTH - 1) + msb_len = escaped.bit_length() - 1 + low_len = escaped - (1 << msb_len) + cw, cl = long_codes[msb_len + _MAX_LONG_LENGTH] + writer.write(cw, cl) + writer.write(low_len, msb_len) + + +def _encode_block(data: bytes, block_index: int) -> bytes: + """Encode one XPRESS9 block: 32-byte header followed by bitstream. + + Implements block assembly (Xpress9EncLz77.c:1367-1512): MTF initial state, two + Huffman tables, the LZ77 token stream, and the 8-word block header with CRC-32C. + """ + tokens = _lz77_tokenize(data) + short_counts, long_counts, _extra_bits = _collect_frequencies(tokens) + + if sum(short_counts) == 0: + short_counts[0] = 1 + if sum(long_counts) == 0: + long_counts[0] = 1 + long_counts[1] = 1 + elif sum(1 for c in long_counts if c > 0) == 1: + for j in range(_LONG_LENGTH_ALPHABET_SIZE): + if long_counts[j] == 0: + long_counts[j] = 1 + break + + short_codes = _build_huffman_codes(short_counts, _SHORT_SYMBOL_ALPHABET_SIZE, _MAX_CODEWORD_LENGTH) + long_codes = _build_huffman_codes(long_counts, _LONG_LENGTH_ALPHABET_SIZE, _MAX_CODEWORD_LENGTH) + + writer = _BitWriter() + + # MTF initial state (Xpress9EncLz77.c:1393-1410). + writer.write(0, 1) # iMtfLastPtr + for _ in range(_ENC_MTF_ENTRY_COUNT): + writer.write(0, 5) # offset 1 encoded as Elias-gamma: msb=0 + + # Write Huffman tables. + short_codes = _encode_huffman_table( + writer, short_codes, short_counts, _SHORT_SYMBOL_ALPHABET_SIZE, _MAX_SHORT_LENGTH + ) + long_codes = _encode_huffman_table(writer, long_codes, long_counts, _LONG_LENGTH_ALPHABET_SIZE, _MAX_SHORT_LENGTH) + + huffman_table_bits = writer.bits_written + _encode_tokens(writer, tokens, short_codes, long_codes) + + exact_bits = writer.bits_written + writer.flush() + bitstream = writer.getvalue() + + comp_size_bits = BLOCK_HEADER_SIZE * 8 + exact_bits + + # Build flags (Xpress9EncLz77.c:1464-1483). + flags = huffman_table_bits & 0x1FFF + flags |= ((_ENC_WINDOW_SIZE_LOG2 - 16) & 7) << 13 + flags |= (_ENC_MTF_ENTRY_COUNT >> 1) << 16 + flags |= ((_ENC_PTR_MIN_MATCH_LENGTH - 3) & 1) << 18 + flags |= ((_ENC_MTF_MIN_MATCH_LENGTH - 2) & 1) << 19 + + header_words = struct.pack( + "<7I", + _XPRESS9_MAGIC, + len(data), + comp_size_bits, + flags, + 0, # reserved + _ENC_SESSION_SIGNATURE, + block_index, + ) + header_crc = _crc32c(header_words) + header = header_words + struct.pack(" bytes: + """Compress data into raw XPRESS9 blocks. + + Encodes the input as a single XPRESS9 block using the fixed "Cosmos Level 6" + parameters (MTF=4, PtrMin=4, MtfMin=2, window 2**16). The result is the raw + block data (32-byte header + bitstream) with no ESE framing. + + Args: + src: The plaintext data to compress. + + Returns: + Raw XPRESS9 block data. + + Raises: + CorruptDataError: If the input is empty. + """ + if not src: + raise CorruptDataError("cannot compress empty input") + return _encode_block(src, 0) diff --git a/dissect/util/compression/lzxpress9_compact.py b/dissect/util/compression/lzxpress9_compact.py new file mode 100644 index 0000000..d84b679 --- /dev/null +++ b/dissect/util/compression/lzxpress9_compact.py @@ -0,0 +1,191 @@ +"""Compact XPRESS9 compression and decompression (ntdll format 0x0005). + +Implements the decompressor and compressor for the undocumented ntdll.dll +compression format 0x0005, introduced in Windows Server 2022 (Build 20348). +Uses the same canonical-Huffman LZ77 engine as the ESE XPRESS9 codec +(:mod:`dissect.util.compression.lzxpress9`) but with a streamlined 10-byte +header instead of XPRESS9's 32-byte block header. + +On-disk layout:: + + bytes 0-3: magic ``0xC039E510`` (LE u32) + bytes 4-5: params (LE u16) -- bits 0-2: window_log index; bit 3: mode + bytes 6-9: control (LE u32) -- bits 0-27: payload bit count (data + 32-bit CRC); + bit 29: compressed flag; bit 31: end-of-stream + bytes 10+: payload (ceil(payload_bits / 8) bytes) + trailing: CRC-32C of original plaintext (32 bits, inside the payload area) + +Reverse-engineered from ``ntdll.dll`` Build 20348 (decompressor at RVA 0x111810). + +Reference implementation: ntcompress (https://github.com/StrongWind1/ntcompress). +""" + +from __future__ import annotations + +import struct +from typing import Final + +from dissect.util.compression.lzxpress9 import ( + _LONG_LENGTH_ALPHABET_SIZE, + _MAX_LONG_LENGTH, + _MAX_MTF, + _MAX_SHORT_LENGTH, + _MAX_SHORT_LENGTH_LOG, + _BitReader, + _CanonicalHuffman, + _copy_match, + _decode_coded_lengths, + _read_match_length, + _take_mtf_offset, +) +from dissect.util.exceptions import CorruptDataError +from dissect.util.hash import crc32c as _crc32c_mod + +_crc32c = _crc32c_mod.crc32c + +MAGIC: Final = 0xC039E510 +"""Block magic -- distinct from XPRESS9's ``0x4E86D72A``.""" + +HEADER_SIZE: Final = 10 +"""Fixed header size: 4 (magic) + 2 (params) + 4 (control).""" + +_WINDOW_LOG_TABLE: Final = (12, 13, 14, 16, 18, 20, 22, 24) + +_PTR_MIN_MATCH: Final = 3 +_MTF_MIN_MATCH: Final = 2 + + +def _short_alphabet_size(window_log: int) -> int: + return (window_log + 16 + _MAX_MTF) << _MAX_SHORT_LENGTH_LOG + + +def _decode_table_2bit(reader: _BitReader, alphabet_size: int) -> _CanonicalHuffman: + """Decode a Huffman table with the compact 2-bit mode prefix. + + Modes: 0 = stored (flat code lengths), 2 = Huffman-coded, 1/3 = error. + """ + mode = reader.read(2) + if mode == 0: + msb = alphabet_size.bit_length() - 1 + short_count = (1 << (msb + 1)) - alphabet_size + return _CanonicalHuffman([msb] * short_count + [msb + 1] * (alphabet_size - short_count)) + if mode == 2: + return _CanonicalHuffman(_decode_coded_lengths(reader, alphabet_size, _MAX_SHORT_LENGTH)) + raise CorruptDataError(f"compact XPRESS9: unsupported table mode {mode}") + + +def decompress(src: bytes, *, verify: bool = True) -> bytes: + """Decompress a compact XPRESS9 stream (ntdll format 0x0005). + + Args: + src: The compressed stream including the 10-byte header. + verify: When True (default), verify the trailing CRC-32C. + + Returns: + The decompressed plaintext. + + Raises: + CorruptDataError: Stream is corrupt, truncated, or CRC mismatch. + """ + if len(src) < HEADER_SIZE: + raise CorruptDataError(f"compact XPRESS9 too short: {len(src)} < {HEADER_SIZE}") + + magic = struct.unpack_from("> 3] + else: + comp_bits = payload_bits - 32 + short_alpha = _short_alphabet_size(window_log) + + reader = _BitReader(payload) + short_table = _decode_table_2bit(reader, short_alpha) + long_table = _decode_table_2bit(reader, _LONG_LENGTH_ALPHABET_SIZE) + + out = bytearray() + mtf: list[int] = [] + last_was_ptr = 0 + + while reader.bits_consumed < comp_bits: + symbol = short_table.decode(reader) + if symbol < 256: + out.append(symbol) + last_was_ptr = 0 + continue + + symbol -= 256 + length = symbol & (_MAX_SHORT_LENGTH - 1) + symbol >>= _MAX_SHORT_LENGTH_LOG + + if length == _MAX_SHORT_LENGTH - 1: + length = _read_match_length(reader, long_table) + + if symbol < _MAX_MTF: + length += _MTF_MIN_MATCH + offset = _take_mtf_offset(mtf, symbol, last_was_ptr=bool(last_was_ptr)) + else: + length += _PTR_MIN_MATCH + msb = symbol - _MAX_MTF + offset = reader.read(msb) + (1 << msb) if msb > 0 else 1 + if len(mtf) < _MAX_MTF: + mtf.insert(0, offset) + else: + mtf.insert(0, offset) + mtf.pop() + + if offset > len(out): + raise CorruptDataError(f"compact XPRESS9 match offset {offset} before start at {len(out)}") + last_was_ptr = 1 + _copy_match(out, offset, length) + + plain = bytes(out) + + if verify: + crc_byte_start = -(-((payload_bits - 32)) // 8) + if crc_byte_start + 4 <= len(payload): + expected = struct.unpack_from(" bytes: + """Compress data into a compact XPRESS9 stream (ntdll format 0x0005). + + Uses uncompressed (literal) mode for simplicity and correctness. The output + is a valid stream that ``RtlDecompressBufferEx(0x0005)`` accepts, though it + does not achieve any size reduction. + + Args: + src: The plaintext to compress. + + Returns: + The compressed stream including header and CRC-32C. + """ + crc = _crc32c(src) + crc_bytes = struct.pack(" int: class Node: + """Huffman tree node.""" + __slots__ = ("children", "is_leaf", "symbol") def __init__(self, symbol: Symbol | None = None, is_leaf: bool = False): @@ -84,6 +91,8 @@ def _build_tree(buf: bytes) -> Node: class BitString: + """LSB-first bitstream reader for LZXPRESS Huffman decoding.""" + def __init__(self): self.source = None self.mask = 0 @@ -91,23 +100,28 @@ def __init__(self): @property def index(self) -> int: + """Current byte offset in the source stream.""" return self.source.tell() def init(self, fh: BinaryIO) -> None: + """Initialize the bitstream from a file-like object.""" self.mask = (_read_16_bit(fh) << 16) + _read_16_bit(fh) self.bits = 32 self.source = fh def read(self, n: int) -> bytes: + """Read n raw bytes from the underlying stream.""" return self.source.read(n) def lookup(self, n: int) -> int: + """Peek at the top n bits without consuming them.""" if n == 0: return 0 return self.mask >> (32 - n) def skip(self, n: int) -> None: + """Consume n bits and refill from the stream.""" self.mask = (self.mask << n) & 0xFFFFFFFF self.bits -= n if self.bits < 16: @@ -115,6 +129,7 @@ def skip(self, n: int) -> None: self.bits += 16 def decode(self, root: Node) -> Symbol: + """Decode one Huffman symbol from the bitstream.""" node = root while not node.is_leaf: bit = self.lookup(1) @@ -123,16 +138,22 @@ def decode(self, root: Node) -> Symbol: return node.symbol -def decompress(src: bytes | BinaryIO) -> bytes: +def decompress(src: bytes | BinaryIO, max_size: int | None = None) -> bytes: """LZXPRESS decompress from a file-like object or bytes. Decompresses until EOF of the input data. Args: src: File-like object or bytes to decompress. + max_size: Optional maximum output size in bytes. When set, raises + ValueError if the decompressed output would exceed this limit. Returns: The decompressed data. + + Raises: + ValueError: If the output exceeds ``max_size``, the Huffman table is + truncated, or the stream contains an invalid match length. """ if not hasattr(src, "read"): src = io.BytesIO(src) @@ -151,11 +172,15 @@ def decompress(src: bytes | BinaryIO) -> bytes: bitstring.init(src) chunk_size = 0 - while chunk_size < 65536 and src.tell() - start_offset < size: + while chunk_size < 65536: symbol = bitstring.decode(root) if symbol < 256: + if max_size is not None and len(dst) >= max_size: + raise ValueError(f"decoded output would exceed the {max_size}-byte limit") dst.append(symbol) chunk_size += 1 + elif symbol == 256: + return bytes(dst) else: symbol -= 256 length = symbol & 0x0F @@ -168,11 +193,18 @@ def decompress(src: bytes | BinaryIO) -> bytes: if length == 270: length = _read_16_bit(bitstring.source) + if length == 0: + length = struct.unpack(" max_size: + raise ValueError(f"decoded output would exceed the {max_size}-byte limit") + remaining = length while remaining > 0: match_size = min(remaining, offset) @@ -182,3 +214,274 @@ def decompress(src: bytes | BinaryIO) -> bytes: chunk_size += length return bytes(dst) + + +# --- Compress ([MS-XCA] S2.1) --- + +_BLOCK_SIZE = 65536 +_TABLE_SIZE = 256 +_SYMBOL_COUNT = 512 +_EOF_SYMBOL = 256 +_MIN_MATCH = 3 +_MAX_CODE_LENGTH = 15 +_MAX_OFFSET = 65535 +_MAX_CHAIN = 64 +_LEN_BYTE_MAX = 0xFF +_LEN_U16_LIMIT = 0x10000 +_MIN_SYMBOLS = 2 + +_Token = tuple[str, int, int] + + +def _high_bit(value: int) -> int: + """Return the index of the highest set bit (GetHighBit, [MS-XCA] S2.1.4.1).""" + return value.bit_length() - 1 + + +def _match_symbol(length: int, distance: int) -> int: + """Return the Huffman symbol for a match (256-511), per [MS-XCA] S2.1.4.1.""" + return _EOF_SYMBOL + min(length - _MIN_MATCH, _MAX_CODE_LENGTH) + 16 * _high_bit(distance) + + +def _lz77(data: bytes) -> list[_Token]: + """Greedy LZ77 factorization with a 3-byte hash chain ([MS-XCA] S2.1.4.1).""" + tokens: list[_Token] = [] + heads: dict[int, int] = {} + chain: list[int] = [-1] * len(data) + n = len(data) + pos = 0 + while pos < n: + best_len = 0 + best_dist = 0 + if pos + _MIN_MATCH <= n: + key = (data[pos] << 16) | (data[pos + 1] << 8) | data[pos + 2] + candidate = heads.get(key, -1) + depth = 0 + while candidate >= 0 and depth < _MAX_CHAIN: + distance = pos - candidate + if distance > _MAX_OFFSET: + break + length = 0 + limit = n - pos + while length < limit and data[candidate + length] == data[pos + length]: + length += 1 + if length > best_len: + best_len = length + best_dist = distance + candidate = chain[candidate] + depth += 1 + if best_len >= _MIN_MATCH and _match_symbol(best_len, best_dist) > _EOF_SYMBOL: + tokens.append(("M", best_len, best_dist)) + advance = best_len + else: + tokens.append(("L", data[pos], 0)) + advance = 1 + end = pos + advance + while pos < end and pos + _MIN_MATCH <= n: + key = (data[pos] << 16) | (data[pos + 1] << 8) | data[pos + 2] + chain[pos] = heads.get(key, -1) + heads[key] = pos + pos += 1 + pos = end + return tokens + + +def _limited_lengths(freqs: Counter[int]) -> dict[int, int]: + """Compute canonical Huffman code lengths capped at 15 bits via package-merge.""" + items = sorted((count, symbol) for symbol, count in freqs.items()) + n = len(items) + lengths = {symbol: 0 for _, symbol in items} + if n == 1: + lengths[items[0][1]] = 1 + return lengths + base = sorted((items[i][0], (i,)) for i in range(n)) + current = list(base) + for _ in range(_MAX_CODE_LENGTH - 1): + packaged = [ + (current[j][0] + current[j + 1][0], current[j][1] + current[j + 1][1]) + for j in range(0, len(current) - 1, 2) + ] + current = sorted(base + packaged) + counts = [0] * n + for _weight, indices in current[: 2 * n - 2]: + for i in indices: + counts[i] += 1 + for i in range(n): + lengths[items[i][1]] = counts[i] + return lengths + + +def _canonical_codes(lengths: dict[int, int]) -> dict[int, int]: + """Assign canonical Huffman codes from bit lengths ([MS-XCA] S2.2.4 ordering).""" + ordered = sorted(symbol for symbol, length in lengths.items() if length > 0) + ordered.sort(key=lambda symbol: (lengths[symbol], symbol)) + codes: dict[int, int] = {} + code = 0 + prev_length = 0 + for symbol in ordered: + code <<= lengths[symbol] - prev_length + codes[symbol] = code + code += 1 + prev_length = lengths[symbol] + return codes + + +def _pack_table(lengths: dict[int, int]) -> bytearray: + """Pack 512 bit lengths into the 256-byte table ([MS-XCA] S2.1.4.3).""" + table = bytearray(_TABLE_SIZE) + for symbol, length in lengths.items(): + if symbol % 2 == 0: + table[symbol // 2] |= length & 0x0F + else: + table[symbol // 2] |= (length & 0x0F) << 4 + return table + + +class _BitWriter: + """Output engine: 16-bit words with two deferred slots ([MS-XCA] S2.1.4.3).""" + + __slots__ = ("buf", "free", "pos", "pos1", "pos2", "word") + + def __init__(self, table: bytearray) -> None: + self.buf = bytearray(table) + self.buf += b"\x00\x00\x00\x00" + self.pos1 = _TABLE_SIZE + self.pos2 = _TABLE_SIZE + 2 + self.pos = _TABLE_SIZE + 4 + self.free = 16 + self.word = 0 + + def _emit_word(self) -> None: + """Write the completed 16-bit word to pos1 and rotate.""" + self.buf[self.pos1] = self.word & 0xFF + self.buf[self.pos1 + 1] = (self.word >> 8) & 0xFF + self.pos1 = self.pos2 + self.pos2 = self.pos + self.buf += b"\x00\x00" + self.pos += 2 + + def write_bits(self, count: int, bits: int) -> None: + """Append count bits (MSB-first) to the deferred bit stream.""" + if count == 0: + return + if self.free >= count: + self.free -= count + self.word = (self.word << count) + bits + else: + self.word = (self.word << self.free) + (bits >> (count - self.free)) + self.free -= count + self._emit_word() + self.free += 16 + self.word = bits + + def write_byte(self, value: int) -> None: + """Write a raw extra-length byte at the current cursor.""" + self.buf.append(value & 0xFF) + self.pos += 1 + + def write_u16(self, value: int) -> None: + """Write a raw little-endian uint16 extra-length value.""" + self.buf += value.to_bytes(2, "little") + self.pos += 2 + + def write_u32(self, value: int) -> None: + """Write a raw little-endian uint32 extra-length value (v10.0 escape).""" + self.buf += value.to_bytes(4, "little") + self.pos += 4 + + def finish(self) -> bytes: + """Flush the pending bits and a trailing zero word.""" + self.word <<= self.free + self.buf[self.pos1] = self.word & 0xFF + self.buf[self.pos1 + 1] = (self.word >> 8) & 0xFF + self.buf[self.pos2] = 0 + self.buf[self.pos2 + 1] = 0 + return bytes(self.buf[: self.pos]) + + +def _write_length_extra(writer: _BitWriter, length: int) -> None: + """Emit extra bytes past the 4-bit match-length nibble ([MS-XCA] S2.1.4.3).""" + value = length - _MIN_MATCH + if value < _MAX_CODE_LENGTH: + return + value -= _MAX_CODE_LENGTH + if value < _LEN_BYTE_MAX: + writer.write_byte(value) + return + writer.write_byte(_LEN_BYTE_MAX) + value += _MAX_CODE_LENGTH + if value < _LEN_U16_LIMIT: + writer.write_u16(value) + else: + writer.write_u16(0) + writer.write_u32(value) + + +def _encode_block(tokens: list[_Token], *, is_last: bool) -> bytes: + """Huffman-encode one block's tokens into a table + bit stream ([MS-XCA] S2.1).""" + freqs: Counter[int] = Counter() + for token in tokens: + if token[0] == "L": + freqs[token[1]] += 1 + else: + freqs[_match_symbol(token[1], token[2])] += 1 + if is_last: + freqs[_EOF_SYMBOL] += 1 + filler = 0 + while len(freqs) < _MIN_SYMBOLS: + if filler not in freqs: + freqs[filler] = 1 + filler += 1 + lengths = _limited_lengths(freqs) + codes = _canonical_codes(lengths) + writer = _BitWriter(_pack_table(lengths)) + for token in tokens: + if token[0] == "L": + writer.write_bits(lengths[token[1]], codes[token[1]]) + continue + length, distance = token[1], token[2] + symbol = _match_symbol(length, distance) + writer.write_bits(lengths[symbol], codes[symbol]) + _write_length_extra(writer, length) + high = _high_bit(distance) + writer.write_bits(high, distance - (1 << high)) + if is_last: + writer.write_bits(lengths[_EOF_SYMBOL], codes[_EOF_SYMBOL]) + return writer.finish() + + +def compress(src: bytes | BinaryIO) -> bytes: + """LZXPRESS Huffman compress from a file-like object or bytes. + + Runs a greedy LZ77 pass, segments the tokens into 64 KiB output blocks each + with its own canonical Huffman table, and appends the EOF symbol after the + final block. Per [MS-XCA] S2.1. + + Args: + src: File-like object or bytes to compress. + + Returns: + The compressed data. + """ + if hasattr(src, "read"): + src = src.read() + + data = bytes(src) + tokens = _lz77(data) + out = bytearray() + index = 0 + produced = 0 + count = len(tokens) + while True: + block: list[_Token] = [] + base = produced + while index < count and produced < base + _BLOCK_SIZE: + token = tokens[index] + block.append(token) + produced += 1 if token[0] == "L" else token[1] + index += 1 + is_last = index >= count + out += _encode_block(block, is_last=is_last) + if is_last: + break + return bytes(out) diff --git a/dissect/util/compression/zlibstream.py b/dissect/util/compression/zlibstream.py new file mode 100644 index 0000000..84268cd --- /dev/null +++ b/dissect/util/compression/zlibstream.py @@ -0,0 +1,38 @@ +"""ZLIB (RFC 1950) compression and decompression. + +Used by ntdll.dll RtlCompressBuffer with COMPRESSION_FORMAT_ZLIB (0x0008). +""" + +from __future__ import annotations + +import zlib as _zlib +from typing import BinaryIO + + +def decompress(src: bytes | BinaryIO) -> bytes: + """Decompress a ZLIB stream (RFC 1950 header + DEFLATE + Adler-32). + + Args: + src: File-like object or bytes to decompress. + + Returns: + The decompressed data. + """ + if hasattr(src, "read"): + src = src.read() + return _zlib.decompress(src) + + +def compress(src: bytes | BinaryIO, level: int = 1) -> bytes: + """Compress data into a ZLIB stream (RFC 1950). + + Args: + src: File-like object or bytes to compress. + level: Compression level 0-9 (default 1). + + Returns: + The compressed ZLIB stream. + """ + if hasattr(src, "read"): + src = src.read() + return _zlib.compress(src, level) diff --git a/dissect/util/hash/__init__.py b/dissect/util/hash/__init__.py index c469368..1114dea 100644 --- a/dissect/util/hash/__init__.py +++ b/dissect/util/hash/__init__.py @@ -1,6 +1,15 @@ +"""Hash utilities providing CRC-32, CRC-32C, CRC-64/NVME, and Jenkins lookup8. + +Selects between a native Rust CRC-32C implementation (when available via +dissect.util._native) and a pure-Python fallback. The native version is +exposed as ``crc32c_native`` and the Python version as ``crc32c_python``; +importing ``crc32c`` gives whichever is available, preferring native. +""" + from __future__ import annotations from dissect.util.hash import crc32c +from dissect.util.hash.crc64 import crc64 crc32c_python = crc32c @@ -26,5 +35,6 @@ "crc32c", "crc32c_native", "crc32c_python", + "crc64", "jenkins", ] diff --git a/dissect/util/hash/crc64.py b/dissect/util/hash/crc64.py new file mode 100644 index 0000000..18c0b8a --- /dev/null +++ b/dissect/util/hash/crc64.py @@ -0,0 +1,104 @@ +"""CRC-64/NVME checksum using reflected polynomial 0x9A6C9329AC4BC9B5.""" + +# fmt: off +from __future__ import annotations + +_TABLE = ( + 0x0000000000000000, 0x7f6ef0c830358979, 0xfedde190606b12f2, 0x81b31158505e9b8b, + 0xc962e5739841b68f, 0xb60c15bba8743ff6, 0x37bf04e3f82aa47d, 0x48d1f42bc81f2d04, + 0xa61cecb46814fe75, 0xd9721c7c5821770c, 0x58c10d24087fec87, 0x27affdec384a65fe, + 0x6f7e09c7f05548fa, 0x1010f90fc060c183, 0x91a3e857903e5a08, 0xeecd189fa00bd371, + 0x78e0ff3b88be6f81, 0x078e0ff3b88be6f8, 0x863d1eabe8d57d73, 0xf953ee63d8e0f40a, + 0xb1821a4810ffd90e, 0xceecea8020ca5077, 0x4f5ffbd87094cbfc, 0x30310b1040a14285, + 0xdefc138fe0aa91f4, 0xa192e347d09f188d, 0x2021f21f80c18306, 0x5f4f02d7b0f40a7f, + 0x179ef6fc78eb277b, 0x68f0063448deae02, 0xe943176c18803589, 0x962de7a428b5bcf0, + 0xf1c1fe77117cdf02, 0x8eaf0ebf2149567b, 0x0f1c1fe77117cdf0, 0x7072ef2f41224489, + 0x38a31b04893d698d, 0x47cdebccb908e0f4, 0xc67efa94e9567b7f, 0xb9100a5cd963f206, + 0x57dd12c379682177, 0x28b3e20b495da80e, 0xa900f35319033385, 0xd66e039b2936bafc, + 0x9ebff7b0e12997f8, 0xe1d10778d11c1e81, 0x606216208142850a, 0x1f0ce6e8b1770c73, + 0x8921014c99c2b083, 0xf64ff184a9f739fa, 0x77fce0dcf9a9a271, 0x08921014c99c2b08, + 0x4043e43f0183060c, 0x3f2d14f731b68f75, 0xbe9e05af61e814fe, 0xc1f0f56751dd9d87, + 0x2f3dedf8f1d64ef6, 0x50531d30c1e3c78f, 0xd1e00c6891bd5c04, 0xae8efca0a188d57d, + 0xe65f088b6997f879, 0x9931f84359a27100, 0x1882e91b09fcea8b, 0x67ec19d339c963f2, + 0xd75adabd7a6e2d6f, 0xa8342a754a5ba416, 0x29873b2d1a053f9d, 0x56e9cbe52a30b6e4, + 0x1e383fcee22f9be0, 0x6156cf06d21a1299, 0xe0e5de5e82448912, 0x9f8b2e96b271006b, + 0x71463609127ad31a, 0x0e28c6c1224f5a63, 0x8f9bd7997211c1e8, 0xf0f5275142244891, + 0xb824d37a8a3b6595, 0xc74a23b2ba0eecec, 0x46f932eaea507767, 0x3997c222da65fe1e, + 0xafba2586f2d042ee, 0xd0d4d54ec2e5cb97, 0x5167c41692bb501c, 0x2e0934dea28ed965, + 0x66d8c0f56a91f461, 0x19b6303d5aa47d18, 0x980521650afae693, 0xe76bd1ad3acf6fea, + 0x09a6c9329ac4bc9b, 0x76c839faaaf135e2, 0xf77b28a2faafae69, 0x8815d86aca9a2710, + 0xc0c42c4102850a14, 0xbfaadc8932b0836d, 0x3e19cdd162ee18e6, 0x41773d1952db919f, + 0x269b24ca6b12f26d, 0x59f5d4025b277b14, 0xd846c55a0b79e09f, 0xa72835923b4c69e6, + 0xeff9c1b9f35344e2, 0x90973171c366cd9b, 0x1124202993385610, 0x6e4ad0e1a30ddf69, + 0x8087c87e03060c18, 0xffe938b633338561, 0x7e5a29ee636d1eea, 0x0134d92653589793, + 0x49e52d0d9b47ba97, 0x368bddc5ab7233ee, 0xb738cc9dfb2ca865, 0xc8563c55cb19211c, + 0x5e7bdbf1e3ac9dec, 0x21152b39d3991495, 0xa0a63a6183c78f1e, 0xdfc8caa9b3f20667, + 0x97193e827bed2b63, 0xe877ce4a4bd8a21a, 0x69c4df121b863991, 0x16aa2fda2bb3b0e8, + 0xf86737458bb86399, 0x8709c78dbb8deae0, 0x06bad6d5ebd3716b, 0x79d4261ddbe6f812, + 0x3105d23613f9d516, 0x4e6b22fe23cc5c6f, 0xcfd833a67392c7e4, 0xb0b6c36e43a74e9d, + 0x9a6c9329ac4bc9b5, 0xe50263e19c7e40cc, 0x64b172b9cc20db47, 0x1bdf8271fc15523e, + 0x530e765a340a7f3a, 0x2c608692043ff643, 0xadd397ca54616dc8, 0xd2bd67026454e4b1, + 0x3c707f9dc45f37c0, 0x431e8f55f46abeb9, 0xc2ad9e0da4342532, 0xbdc36ec59401ac4b, + 0xf5129aee5c1e814f, 0x8a7c6a266c2b0836, 0x0bcf7b7e3c7593bd, 0x74a18bb60c401ac4, + 0xe28c6c1224f5a634, 0x9de29cda14c02f4d, 0x1c518d82449eb4c6, 0x633f7d4a74ab3dbf, + 0x2bee8961bcb410bb, 0x548079a98c8199c2, 0xd53368f1dcdf0249, 0xaa5d9839ecea8b30, + 0x449080a64ce15841, 0x3bfe706e7cd4d138, 0xba4d61362c8a4ab3, 0xc52391fe1cbfc3ca, + 0x8df265d5d4a0eece, 0xf29c951de49567b7, 0x732f8445b4cbfc3c, 0x0c41748d84fe7545, + 0x6bad6d5ebd3716b7, 0x14c39d968d029fce, 0x95708ccedd5c0445, 0xea1e7c06ed698d3c, + 0xa2cf882d2576a038, 0xdda178e515432941, 0x5c1269bd451db2ca, 0x237c997575283bb3, + 0xcdb181ead523e8c2, 0xb2df7122e51661bb, 0x336c607ab548fa30, 0x4c0290b2857d7349, + 0x04d364994d625e4d, 0x7bbd94517d57d734, 0xfa0e85092d094cbf, 0x856075c11d3cc5c6, + 0x134d926535897936, 0x6c2362ad05bcf04f, 0xed9073f555e26bc4, 0x92fe833d65d7e2bd, + 0xda2f7716adc8cfb9, 0xa54187de9dfd46c0, 0x24f29686cda3dd4b, 0x5b9c664efd965432, + 0xb5517ed15d9d8743, 0xca3f8e196da80e3a, 0x4b8c9f413df695b1, 0x34e26f890dc31cc8, + 0x7c339ba2c5dc31cc, 0x035d6b6af5e9b8b5, 0x82ee7a32a5b7233e, 0xfd808afa9582aa47, + 0x4d364994d625e4da, 0x3258b95ce6106da3, 0xb3eba804b64ef628, 0xcc8558cc867b7f51, + 0x8454ace74e645255, 0xfb3a5c2f7e51db2c, 0x7a894d772e0f40a7, 0x05e7bdbf1e3ac9de, + 0xeb2aa520be311aaf, 0x944455e88e0493d6, 0x15f744b0de5a085d, 0x6a99b478ee6f8124, + 0x224840532670ac20, 0x5d26b09b16452559, 0xdc95a1c3461bbed2, 0xa3fb510b762e37ab, + 0x35d6b6af5e9b8b5b, 0x4ab846676eae0222, 0xcb0b573f3ef099a9, 0xb465a7f70ec510d0, + 0xfcb453dcc6da3dd4, 0x83daa314f6efb4ad, 0x0269b24ca6b12f26, 0x7d0742849684a65f, + 0x93ca5a1b368f752e, 0xeca4aad306bafc57, 0x6d17bb8b56e467dc, 0x12794b4366d1eea5, + 0x5aa8bf68aecec3a1, 0x25c64fa09efb4ad8, 0xa4755ef8cea5d153, 0xdb1bae30fe90582a, + 0xbcf7b7e3c7593bd8, 0xc399472bf76cb2a1, 0x422a5673a732292a, 0x3d44a6bb9707a053, + 0x759552905f188d57, 0x0afba2586f2d042e, 0x8b48b3003f739fa5, 0xf42643c80f4616dc, + 0x1aeb5b57af4dc5ad, 0x6585ab9f9f784cd4, 0xe436bac7cf26d75f, 0x9b584a0fff135e26, + 0xd389be24370c7322, 0xace74eec0739fa5b, 0x2d545fb4576761d0, 0x523aaf7c6752e8a9, + 0xc41748d84fe75459, 0xbb79b8107fd2dd20, 0x3acaa9482f8c46ab, 0x45a459801fb9cfd2, + 0x0d75adabd7a6e2d6, 0x721b5d63e7936baf, 0xf3a84c3bb7cdf024, 0x8cc6bcf387f8795d, + 0x620ba46c27f3aa2c, 0x1d6554a417c62355, 0x9cd645fc4798b8de, 0xe3b8b53477ad31a7, + 0xab69411fbfb21ca3, 0xd407b1d78f8795da, 0x55b4a08fdfd90e51, 0x2ada5047efec8728, +) +# fmt: on + +_MASK = 0xFFFFFFFFFFFFFFFF + + +def update(crc: int, data: bytes) -> int: + """Update CRC-64/NVME checksum with data. + + Args: + crc: The initial value of the checksum. + data: The data to update the checksum with. + + Returns: + The updated CRC-64 checksum value. + """ + crc = crc ^ _MASK + for b in data: + table_idx = (crc ^ b) & 0xFF + crc = _TABLE[table_idx] ^ ((crc >> 8) & _MASK) + return crc ^ _MASK + + +def crc64(data: bytes, value: int = 0) -> int: + """Calculate CRC-64/NVME checksum of some data, with an optional initial value. + + Args: + data: The data to calculate the checksum of. + value: The initial value of the checksum. Default is 0. + + Returns: + The CRC-64 checksum. + """ + return update(value, data) & _MASK diff --git a/tests/compression/test_deflate.py b/tests/compression/test_deflate.py new file mode 100644 index 0000000..ab775fd --- /dev/null +++ b/tests/compression/test_deflate.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import hashlib +import zlib +from typing import TYPE_CHECKING + +import pytest + +from dissect.util.compression import deflate + +if TYPE_CHECKING: + from pytest_benchmark.fixture import BenchmarkFixture + + +# Raw DEFLATE of the 4096-byte A-Z pattern, from ntdll RtlCompressBuffer +# (COMPRESSION_FORMAT_DEFLATE, 0x0007) on Win11 Build 26100. +GOLD = ( + "737276717573f7f0f4f2f6f1f5f30f080c0a0e090d0b8f888c72741a95190d83d174309a1746cb83d13271b45e18ad1b47db07a3" + "6da4d176e2685b79b4bf30da331aed190d939e1100" +) +GOLD_DIGEST = "45508be868dafa697653cbc4cc934b0ce5266eb7e87a5a9ac8fc8005ecc03189" + +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = "2bc94855282ccd4cce56482aca2fcf5348cbaf50c82acd2d2856c82f4b2d5228014ae72456552aa4e4a7eb8179a38ac90a0d00" + + +def test_deflate_decompress() -> None: + assert hashlib.sha256(deflate.decompress(bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +def test_deflate_compress() -> None: + assert deflate.decompress(deflate.compress(INPUT)) == INPUT + assert deflate.compress(INPUT) == bytes.fromhex(COMPRESSED) + + +def test_deflate_decompress_max_size() -> None: + compressed = deflate.compress(INPUT) + assert deflate.decompress(compressed, max_size=len(INPUT)) == INPUT + + with pytest.raises(zlib.error): + deflate.decompress(compressed, max_size=10) + + +@pytest.mark.benchmark +def test_benchmark_deflate_decompress(benchmark: BenchmarkFixture) -> None: + assert hashlib.sha256(benchmark(deflate.decompress, bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +@pytest.mark.benchmark +def test_benchmark_deflate_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(deflate.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_lz4.py b/tests/compression/test_lz4.py index 5220cf6..4830b95 100644 --- a/tests/compression/test_lz4.py +++ b/tests/compression/test_lz4.py @@ -43,12 +43,34 @@ ) +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "f01074686520717569636b2062726f776e20666f78206a756d7073206f766572201f00916c617a7920646f672e0e000f2d00ff2050" + "646f672e20" +) + + @pytest.mark.parametrize(*PARAMS) def test_lz4_decompress(lz4: ModuleType, data: str, digest: str) -> None: assert hashlib.sha256(lz4.decompress(bytes.fromhex(data))).hexdigest() == digest +def test_lz4_compress() -> None: + # Only the pure-Python codec implements compress(); the native one is decompress-only. + from dissect.util.compression import lz4_python + + assert lz4_python.decompress(lz4_python.compress(INPUT), len(INPUT)) == INPUT + assert lz4_python.compress(INPUT) == bytes.fromhex(COMPRESSED) + + @pytest.mark.benchmark @pytest.mark.parametrize(*PARAMS) def test_benchmark_lz4_decompress(lz4: ModuleType, data: str, digest: str, benchmark: BenchmarkFixture) -> None: assert hashlib.sha256(benchmark(lz4.decompress, bytes.fromhex(data))).hexdigest() == digest + + +@pytest.mark.benchmark +def test_benchmark_lz4_compress(benchmark: BenchmarkFixture) -> None: + from dissect.util.compression import lz4_python + + assert benchmark(lz4_python.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_lznt1.py b/tests/compression/test_lznt1.py index 9dca962..31b32e5 100644 --- a/tests/compression/test_lznt1.py +++ b/tests/compression/test_lznt1.py @@ -132,12 +132,43 @@ ) +# LZNT1 of the 4096-byte A-Z pattern, from ntdll RtlCompressBuffer +# (COMPRESSION_FORMAT_LZNT1, 0x0002) on Win11 Build 26100. +GOLD = ( + "0fb1004243444546474849004a4b4c4d4e4f5051005253545556575859fc5a41ffcf5f80ff819f839f833f85ffdf86df867f881f8abf8bbf" + "8b5f8dff8eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0e" + "ff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e" + "9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f" + "0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0eff9f0e9f0e9f0e9f0e9f0e9f0e9f0e9f0e0f9f0e9f0e9f0e910e" +) +GOLD_DIGEST = "45508be868dafa697653cbc4cc934b0ce5266eb7e87a5a9ac8fc8005ecc03189" + +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "33b0007468652071756963006b2062726f776e2000666f78206a756d708073206f7665722001f0006c617a7920646f67062e023434b1" +) + + @pytest.mark.parametrize(*PARAMS) def test_lznt1_decompress(data: str, digest: str) -> None: assert hashlib.sha256(lznt1.decompress(bytes.fromhex(data)).rstrip(b"\x00")).hexdigest() == digest +def test_lznt1_decompress_gold() -> None: + assert hashlib.sha256(lznt1.decompress(bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +def test_lznt1_compress() -> None: + assert lznt1.decompress(lznt1.compress(INPUT)) == INPUT + assert lznt1.compress(INPUT) == bytes.fromhex(COMPRESSED) + + @pytest.mark.benchmark @pytest.mark.parametrize(*PARAMS) def test_benchmark_lznt1_decompress(data: str, digest: str, benchmark: BenchmarkFixture) -> None: assert hashlib.sha256(benchmark(lznt1.decompress, bytes.fromhex(data)).rstrip(b"\x00")).hexdigest() == digest + + +@pytest.mark.benchmark +def test_benchmark_lznt1_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(lznt1.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_lzxpress.py b/tests/compression/test_lzxpress.py index 7bd0ea2..9939d6f 100644 --- a/tests/compression/test_lzxpress.py +++ b/tests/compression/test_lzxpress.py @@ -42,12 +42,37 @@ ) +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "0100000074686520717569636b2062726f776e20666f78206a756d7073206f76657220f100ffff7f006c617a7920646f672e6a00" + "67010fff3401" +) + + @pytest.mark.parametrize(*PARAMS) def test_lzxpress_decompress(data: str, digest: str) -> None: assert hashlib.sha256(lzxpress.decompress(bytes.fromhex(data))).hexdigest() == digest +def test_lzxpress_compress() -> None: + assert lzxpress.decompress(lzxpress.compress(INPUT)) == INPUT + assert lzxpress.compress(INPUT) == bytes.fromhex(COMPRESSED) + + +def test_lzxpress_decompress_max_size() -> None: + compressed = lzxpress.compress(INPUT) + assert lzxpress.decompress(compressed, max_size=len(INPUT)) == INPUT + + with pytest.raises(ValueError, match="max_size"): + lzxpress.decompress(compressed, max_size=10) + + @pytest.mark.benchmark @pytest.mark.parametrize(*PARAMS) def test_benchmark_lzxpress_decompress(data: str, digest: str, benchmark: BenchmarkFixture) -> None: assert hashlib.sha256(benchmark(lzxpress.decompress, bytes.fromhex(data))).hexdigest() == digest + + +@pytest.mark.benchmark +def test_benchmark_lzxpress_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(lzxpress.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_lzxpress9.py b/tests/compression/test_lzxpress9.py new file mode 100644 index 0000000..8346ecb --- /dev/null +++ b/tests/compression/test_lzxpress9.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +import pytest + +from dissect.util.compression import lzxpress9 + +if TYPE_CHECKING: + from pytest_benchmark.fixture import BenchmarkFixture + + +# Raw XPRESS9 blocks (ESE 5-byte record header stripped) from the MIT ESE C +# reference encoder, the same code esent.dll links. Session signature is +# normalized to 0x12345678; the decoder accepts any valid signature. +PARAMS = ( + ("data", "size", "digest"), + [ + pytest.param( + "2ad7864e68010000d00200001b00060000000000eeadd4ba0000000015cc7f96000000e0c28229028e5c5932668d801127f6" + "dcd92160c69e0702565cd972e08c803d37a69c107061c114011b86bc782260c29e39ba1addfe6d6f", + 360, + "0c698b10ca43e8c4ad7225fc3fa3df99373cf62c4831b702f6e1f2b5d3601683", + id="fox", + ), + pytest.param( + "2ad7864ee8030000470100001b0006000000000043718bbc0000000012e14ffa000000000020fca33b", + 1000, + "541b3e9daa09b20bf85fa273e5cbd3e80185aa4ec298e765db87742b70138a53", + id="zeros", + ), + pytest.param( + "2ad7864e480e0000030400004101060000000000cea7f001000000005ec0066e000020c266a6000040dcc56e12be7b5be118" + "5559d57fe7dbc9d6e475a63af30600701fc17d040000466aa415fe0e710c9f016bf88443d64e48123a775f682f68b42d5945" + "83b84167ea52bd351abefc4512b547cac437c6eb197fd5edfe789e9700", + 3656, + "98aa38df85425a9709e9cdb479eddeb556a20583403681eb198d2ee82184c5ad", + id="paragraph", + ), + ], +) + +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "2ad7864e68010000cd0200001b000600000000007856341200000000c140bd45000000e0c28229028e5c5932668d801127f6dcd9" + "2160c69e0702565cd972e08c803d37a69c107061c114011b86bc782260c29e393a04eddf050d" +) + + +@pytest.mark.parametrize(*PARAMS) +def test_lzxpress9_decompress(data: str, size: int, digest: str) -> None: + assert hashlib.sha256(lzxpress9.decompress(bytes.fromhex(data))).hexdigest() == digest + + +@pytest.mark.parametrize(*PARAMS) +def test_lzxpress9_decompressed_size(data: str, size: int, digest: str) -> None: + assert lzxpress9.decompressed_size(bytes.fromhex(data)) == size + + +def test_lzxpress9_compress() -> None: + assert lzxpress9.decompress(lzxpress9.compress(INPUT)) == INPUT + assert lzxpress9.compress(INPUT) == bytes.fromhex(COMPRESSED) + + +@pytest.mark.benchmark +@pytest.mark.parametrize(*PARAMS) +def test_benchmark_lzxpress9_decompress(data: str, size: int, digest: str, benchmark: BenchmarkFixture) -> None: + assert hashlib.sha256(benchmark(lzxpress9.decompress, bytes.fromhex(data))).hexdigest() == digest + + +@pytest.mark.benchmark +def test_benchmark_lzxpress9_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(lzxpress9.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_lzxpress9_compact.py b/tests/compression/test_lzxpress9_compact.py new file mode 100644 index 0000000..ce90892 --- /dev/null +++ b/tests/compression/test_lzxpress9_compact.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +import pytest + +from dissect.util.compression import lzxpress9_compact +from dissect.util.exceptions import CorruptDataError + +if TYPE_CHECKING: + from pytest_benchmark.fixture import BenchmarkFixture + + +# Raw compact XPRESS9 streams from ntdll RtlCompressBuffer(0x0005) on +# Server 2022 Build 20348. All decompress to the same 4096-byte A-Z pattern. +PARAMS = ( + ("data", "digest"), + [ + pytest.param( + "10e539c007404a0000a000a0f86783038941f998", + "ad7facb2586fc6e966c004d7d1d16b024f5805ff7cb47c7a85dabd8b48892ca7", + id="zeros_4096", + ), + pytest.param( + "10e539c007404a0000a040b0f8678303e9517205", + "6896d9ea3f73a4434f5832bc65714e7d066f177373f36f34dc8a6f735daa41b1", + id="allA_4096", + ), + pytest.param( + "10e539c007403f0000a000a0785a67ebc803", + "f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b", + id="zeros_64", + ), + pytest.param( + "10e539c00740540000a0a08a6af9cf0507021fea46", + "116109c6e03f7d2cef4c67aff339970054357a2c75e3bb90f9adc388c0a1ffdd", + id="alt_AA55_4096", + ), + ], +) + +# Gold vector: 4096-byte A-Z pattern compressed with format 0x0005. +# Byte-identical across Server 2022 (Build 20348) and Server 2025 (Build 26100). +GOLD = ( + "10e539c007402f0100a040883011a2c4889320498a3419b2e4c853a048893215" + "aad4a8d3a0498b60f3cfed56dc7a8d3d" +) +GOLD_DIGEST = "45508be868dafa697653cbc4cc934b0ce5266eb7e87a5a9ac8fc8005ecc03189" + +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 + + +@pytest.mark.parametrize(*PARAMS) +def test_lzxpress9_compact_decompress(data: str, digest: str) -> None: + assert hashlib.sha256(lzxpress9_compact.decompress(bytes.fromhex(data))).hexdigest() == digest + + +def test_lzxpress9_compact_decompress_gold() -> None: + assert hashlib.sha256(lzxpress9_compact.decompress(bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +def test_lzxpress9_compact_decompress_uncompressed() -> None: + comp = bytes.fromhex("10e539c00740280000800051537d52") + assert lzxpress9_compact.decompress(comp) == b"\x00" + + +def test_lzxpress9_compact_decompress_skip_verify() -> None: + assert hashlib.sha256(lzxpress9_compact.decompress(bytes.fromhex(GOLD), verify=False)).hexdigest() == GOLD_DIGEST + + +def test_lzxpress9_compact_decompress_bad_magic() -> None: + bad = b"\x00\x00\x00\x00" + bytes.fromhex(GOLD)[4:] + with pytest.raises(CorruptDataError, match="magic"): + lzxpress9_compact.decompress(bad) + + +def test_lzxpress9_compact_decompress_truncated() -> None: + with pytest.raises(CorruptDataError, match="too short"): + lzxpress9_compact.decompress(b"\x10\xe5\x39\xc0") + + +def test_lzxpress9_compact_compress() -> None: + assert lzxpress9_compact.decompress(lzxpress9_compact.compress(INPUT)) == INPUT + + +@pytest.mark.benchmark +def test_benchmark_lzxpress9_compact_decompress(benchmark: BenchmarkFixture) -> None: + assert hashlib.sha256(benchmark(lzxpress9_compact.decompress, bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +@pytest.mark.benchmark +def test_benchmark_lzxpress9_compact_compress(benchmark: BenchmarkFixture) -> None: + result = benchmark(lzxpress9_compact.compress, INPUT) + assert lzxpress9_compact.decompress(result) == INPUT diff --git a/tests/compression/test_lzxpress_huffman.py b/tests/compression/test_lzxpress_huffman.py index b223454..6703388 100644 --- a/tests/compression/test_lzxpress_huffman.py +++ b/tests/compression/test_lzxpress_huffman.py @@ -51,12 +51,55 @@ ) +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "0000000000000000000000000000000003000000000000060000000000000000000000000000000000000000000000006066466666666636" + "6664465555050000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000005000000000000000000000000000000000000000000000000050000000000005000000000000000" + "0000000000000050000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000" + "000000000000000000000000000000000000000000000000000000000000000054ffb61ef0d85230c6fe0366eb7681cf2872f3afa27c3a32" + "4d5dcc660000ff3401" +) + + @pytest.mark.parametrize(*PARAMS) def test_lzxpress_huffman_decompress(data: str, digest: str) -> None: assert hashlib.sha256(lzxpress_huffman.decompress(bytes.fromhex(data))).hexdigest() == digest +def test_lzxpress_huffman_compress() -> None: + assert lzxpress_huffman.decompress(lzxpress_huffman.compress(INPUT)) == INPUT + assert lzxpress_huffman.compress(INPUT) == bytes.fromhex(COMPRESSED) + + +def test_lzxpress_huffman_decompress_max_size() -> None: + compressed = lzxpress_huffman.compress(INPUT) + assert lzxpress_huffman.decompress(compressed, max_size=len(INPUT)) == INPUT + + with pytest.raises(ValueError, match="limit"): + lzxpress_huffman.decompress(compressed, max_size=10) + + +@pytest.mark.parametrize("data", [b"A", b"AB", b"hello", b"the quick brown fox"]) +def test_lzxpress_huffman_eof_symbol(data: bytes) -> None: + # Regression: the end-of-block symbol (256) must terminate decoding instead of + # being decoded as a match, which appended trailing garbage for short streams. + assert lzxpress_huffman.decompress(lzxpress_huffman.compress(data)) == data + + +def test_lzxpress_huffman_long_match() -> None: + # Regression: a match longer than 0xFFFF bytes uses the v10.0 uint16(0)+uint32 + # length escape, which the decoder previously did not read. + data = b"A" * 70000 + assert lzxpress_huffman.decompress(lzxpress_huffman.compress(data)) == data + + @pytest.mark.benchmark @pytest.mark.parametrize(*PARAMS) def test_benchmark_lzxpress_huffman_decompress(data: str, digest: str, benchmark: BenchmarkFixture) -> None: assert hashlib.sha256(benchmark(lzxpress_huffman.decompress, bytes.fromhex(data))).hexdigest() == digest + + +@pytest.mark.benchmark +def test_benchmark_lzxpress_huffman_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(lzxpress_huffman.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/compression/test_zlibstream.py b/tests/compression/test_zlibstream.py new file mode 100644 index 0000000..7b759aa --- /dev/null +++ b/tests/compression/test_zlibstream.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +import pytest + +from dissect.util.compression import zlibstream + +if TYPE_CHECKING: + from pytest_benchmark.fixture import BenchmarkFixture + + +# ZLIB stream of the 4096-byte A-Z pattern, from ntdll RtlCompressBuffer +# (COMPRESSION_FORMAT_ZLIB, 0x0008) on Win11 Build 26100. +GOLD = ( + "7801737276717573f7f0f4f2f6f1f5f30f080c0a0e090d0b8f888c72741a95190d83d174309a1746cb83d13271b45e18ad1b47db" + "07a36da4d176e2685b79b4bf30da331aed190d939e110004d2d7f7" +) +GOLD_DIGEST = "45508be868dafa697653cbc4cc934b0ce5266eb7e87a5a9ac8fc8005ecc03189" + +INPUT = b"the quick brown fox jumps over the lazy dog. " * 8 +COMPRESSED = ( + "78012bc94855282ccd4cce56482aca2fcf5348cbaf50c82acd2d2856c82f4b2d5228014ae72456552aa4e4a7eb8179a38ac90a0d002fc08239" +) + + +def test_zlibstream_decompress() -> None: + assert hashlib.sha256(zlibstream.decompress(bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +def test_zlibstream_compress() -> None: + assert zlibstream.decompress(zlibstream.compress(INPUT)) == INPUT + assert zlibstream.compress(INPUT) == bytes.fromhex(COMPRESSED) + + +@pytest.mark.benchmark +def test_benchmark_zlibstream_decompress(benchmark: BenchmarkFixture) -> None: + assert hashlib.sha256(benchmark(zlibstream.decompress, bytes.fromhex(GOLD))).hexdigest() == GOLD_DIGEST + + +@pytest.mark.benchmark +def test_benchmark_zlibstream_compress(benchmark: BenchmarkFixture) -> None: + assert benchmark(zlibstream.compress, INPUT) == bytes.fromhex(COMPRESSED) diff --git a/tests/test_hash.py b/tests/test_hash.py index 3a57ddb..84e149d 100644 --- a/tests/test_hash.py +++ b/tests/test_hash.py @@ -4,7 +4,7 @@ import pytest -from dissect.util.hash import crc32, jenkins +from dissect.util.hash import crc32, crc64, jenkins if TYPE_CHECKING: from types import ModuleType @@ -45,6 +45,28 @@ def test_crc32c_benchmark(crc32c: ModuleType, benchmark: BenchmarkFixture) -> No benchmark(crc32c.crc32c, b"hello, world!", 0) +@pytest.mark.parametrize( + ("data", "value", "expected"), + [ + # CRC-64/NVME check value (reveng.sourceforge.io CRC catalogue) + (b"123456789", 0, 0xAE8B14860A799888), + (b"", 0, 0), + (b"hello, world!", 0, 0xF8046E40C403F1D0), + (b"hello, world!", 0x12345678, 0x080A99F5C71E0576), + (b"\x00" * 32, 0, 0xCF3473434D4ECF3B), + (b"\xff" * 32, 0, 0xA0A06974C34D63C4), + (bytes(range(32)), 0, 0xB9D9D4A8492CBD7F), + ], +) +def test_crc64(data: bytes, value: int, expected: int) -> None: + assert crc64(data, value) == expected + + +@pytest.mark.benchmark +def test_crc64_benchmark(benchmark: BenchmarkFixture) -> None: + benchmark(crc64, b"hello, world!", 0) + + def test_lookup8_remainder() -> None: ip = b"192.168.1.109" volume = b"/home/roel/nfstest"