Add XZ (LZMA) checksum repair utility

dissect/util/compression/xz.py

@@ -0,0 +1,101 @@
+import io
+from binascii import crc32
+from typing import BinaryIO
+
+from dissect.util.stream import OverlayStream
+
+
+def repair_checksum(fh: BinaryIO) -> BinaryIO:
+    """Repair CRC32 checksums for all headers in an XZ stream.
+
+    FortiOS XZ files have (on purpose) corrupt streams which they read using a modified ``xz`` binary.
+    The only thing changed are the CRC32 checksums, so partially parse the XZ file and fix all of them.
+
+    References:
+        - https://tukaani.org/xz/xz-file-format-1.1.0.txt
+        - https://github.com/Rogdham/python-xz
+
+    Args:
+        fh: A file-like object of an LZMA stream to repair.
+    """
+    size = fh.seek(0, io.SEEK_END)
+    repaired = OverlayStream(fh, size)
+    fh.seek(0)
+
+    header = fh.read(12)
+    # Check header magic
+    if header[:6] != b"\xfd7zXZ\x00":
+        raise ValueError("Not an XZ file")
+
+    # Add correct header CRC32
+    repaired.add(8, _crc32(header[6:8]))
+
+    fh.seek(-12, io.SEEK_END)
+    footer = fh.read(12)
+
+    # Check footer magic
+    if footer[10:12] != b"YZ":
+        raise ValueError("Not an XZ file")
+
+    # Add correct footer CRC32
+    repaired.add(fh.tell() - 12, _crc32(footer[4:10]))
+
+    backward_size = (int.from_bytes(footer[4:8], "little") + 1) * 4
+    fh.seek(-12 - backward_size, io.SEEK_END)
+    index = fh.read(backward_size)
+
+    # Add correct index CRC32
+    repaired.add(fh.tell() - 4, _crc32(index[:-4]))
+
+    # Parse the index
+    isize, nb_records = _mbi(index[1:])
+    index = index[1 + isize : -4]
+    records = []
+    for _ in range(nb_records):
+        if not index:
+            raise ValueError("index size")
+
+        isize, unpadded_size = _mbi(index)
+        if not unpadded_size:
+            raise ValueError("index record unpadded size")
+
+        index = index[isize:]
+        if not index:
+            raise ValueError("index size")
+
+        isize, uncompressed_size = _mbi(index)
+        if not uncompressed_size:
+            raise ValueError("index record uncompressed size")
+
+        index = index[isize:]
+        records.append((unpadded_size, uncompressed_size))
+
+    block_start = size - 12 - backward_size
+    blocks_len = sum((unpadded_size + 3) & ~3 for unpadded_size, _ in records)
+    block_start -= blocks_len
+
+    # Iterate over all blocks and add the correct block header CRC32
+    for unpadded_size, _ in records:
+        fh.seek(block_start)
+
+        block_header = fh.read(1)
+        block_header_size = (block_header[0] + 1) * 4
+        block_header += fh.read(block_header_size - 1)
+        repaired.add(fh.tell() - 4, _crc32(block_header[:-4]))
+
+        block_start += (unpadded_size + 3) & ~3
+
+    return repaired
+
+
+def _mbi(data: bytes) -> tuple[int, int]:
+    value = 0
+    for size, byte in enumerate(data):
+        value |= (byte & 0x7F) << (size * 7)
+        if not byte & 0x80:
+            return size + 1, value
+    raise ValueError("Invalid mbi")
+
+
+def _crc32(data: bytes) -> bytes:
+    return int.to_bytes(crc32(data), 4, "little")

tests/test_compression.py

@@ -1,4 +1,6 @@
 import hashlib
+import lzma
+from io import BytesIO
 
 from dissect.util.compression import (
     lz4,
@@ -7,6 +9,7 @@
     lzxpress,
     lzxpress_huffman,
     sevenbit,
+    xz,
 )
 
 
@@ -254,3 +257,16 @@ def test_sevenbit_decompress_wide():
     result = sevenbit.decompress(bytes.fromhex("b796384d078ddf6db8bc3c9fa7df6e10bd3ca783e67479da7d06"), wide=True)
     target = "7-bit compression test string".encode("utf-16-le")
     assert result == target
+
+
+def test_xz_repair_checksum():
+    buf = BytesIO(
+        bytes.fromhex(
+            "fd377a585a000004deadbeef0200210116000000deadbeefe00fff001e5d003a"
+            "194ace2b0f238ce989a29cfeb182a4e814985366b771770233ca314836000000"
+            "2972e8fd62b18ee300013a8020000000deadbeefdeadbeef020000000004595a"
+        )
+    )
+    repaired = xz.repair_checksum(buf)
+
+    assert lzma.decompress(repaired.read()) == b"test" * 1024