Add compression and hash algorithm support

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",
 ]

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()

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

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("<H", word)
+            group_flags |= 1 << group_count
+        else:
+            length = 1
+            group_items.append(chunk[pos])
+
+        for indexed in range(pos, min(pos + length, len(chunk) - _MIN_MATCH + 1)):
+            table.setdefault(chunk[indexed : indexed + _MIN_MATCH], []).append(indexed)
+
+        pos += length
+        group_count += 1
+
+        if group_count == 8:
+            body.append(group_flags)
+            body += group_items
+            group_flags = 0
+            group_items.clear()
+            group_count = 0
+
+    if group_count:
+        body.append(group_flags)
+        body += group_items
+
+    return bytes(body)
+
+
+def compress(src: bytes | BinaryIO) -> 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("<H", header)
+            result += body
+        else:
+            header = SIGNATURE_MASK | (len(chunk) + 2 - 3)
+            result += struct.pack("<H", header)
+            result += chunk
+
+    return bytes(result)