#!/usr/bin/env python3
"""
make_pe.py - build a small, valid, benign PE32 test fixture. Stdlib only.

The binary imports a handful of Windows APIs that triage heuristics care
about (injection, network, anti-debug, persistence) and embeds indicator
strings, then simply calls ExitProcess(0). It is deterministic (fixed
timestamp) so hashes are stable across runs.

Usage:
    python3 make_pe.py out/invoice_2024.exe
"""

import struct
import sys

IMAGE_BASE = 0x400000
SEC_ALIGN, FILE_ALIGN = 0x1000, 0x200
TIMESTAMP = 1704067200  # 2024-01-01T00:00:00Z

IMPORTS = {
    b"KERNEL32.dll": [b"ExitProcess", b"VirtualAllocEx", b"WriteProcessMemory", b"CreateRemoteThread",
                      b"OpenProcess", b"IsDebuggerPresent", b"GetTempPathA", b"CreateProcessA"],
    b"WININET.dll": [b"InternetOpenA", b"InternetOpenUrlA", b"InternetReadFile"],
    b"ADVAPI32.dll": [b"RegCreateKeyExA", b"RegSetValueExA"],
}

ASCII_STRINGS = [
    b"http://update-cdn.example.test/gate.php",
    b"Mozilla/4.0 (compatible; MSIE 6.0; TestUA/1.0)",
    b"Global\\TestFixtureMutex_9f2a",
    b"Software\\Microsoft\\Windows\\CurrentVersion\\Run",
    b"WinUpdate",
    b"C:\\Users\\Public\\svc.exe",
    b"203.0.113.5",
    b"D:\\build\\fixture\\Release\\invoice.pdb",
    b"BENIGN TEST FIXTURE - this program only calls ExitProcess(0)",
]
WIDE_STRINGS = ["cmd.exe /c ping -n 3 127.0.0.1 & del \"%s\"", "\\\\.\\pipe\\testfixture"]


def align(n, a):
    return (n + a - 1) // a * a


def build_rdata(rdata_rva):
    """Return (bytes, iat_rva_of_ExitProcess, import_dir_rva, import_dir_size)."""
    buf = bytearray()
    n_dlls = len(IMPORTS)
    idt_size = 20 * (n_dlls + 1)
    buf += b"\0" * idt_size  # import directory table, patched later
    thunks = {}  # dll -> (ilt_off, iat_off)
    for dll, funcs in IMPORTS.items():
        ilt_off = len(buf); buf += b"\0" * (4 * (len(funcs) + 1))
        iat_off = len(buf); buf += b"\0" * (4 * (len(funcs) + 1))
        thunks[dll] = (ilt_off, iat_off)
    name_offs = {}
    for dll, funcs in IMPORTS.items():
        for i, f in enumerate(funcs):
            if len(buf) % 2: buf += b"\0"
            name_offs[(dll, f)] = len(buf)
            buf += struct.pack("<H", 0) + f + b"\0"
    dll_offs = {}
    for dll in IMPORTS:
        dll_offs[dll] = len(buf); buf += dll + b"\0"
    # patch thunks and IDT
    exit_iat = None
    for idx, (dll, funcs) in enumerate(IMPORTS.items()):
        ilt_off, iat_off = thunks[dll]
        for i, f in enumerate(funcs):
            rva = rdata_rva + name_offs[(dll, f)]
            struct.pack_into("<I", buf, ilt_off + 4 * i, rva)
            struct.pack_into("<I", buf, iat_off + 4 * i, rva)
            if f == b"ExitProcess":
                exit_iat = rdata_rva + iat_off + 4 * i
        struct.pack_into("<IIIII", buf, 20 * idx, rdata_rva + ilt_off, 0, 0, rdata_rva + dll_offs[dll], rdata_rva + iat_off)
    # strings
    buf += b"\0" * 16
    for s in ASCII_STRINGS:
        buf += s + b"\0"
    buf += b"\0" * 4
    for s in WIDE_STRINGS:
        buf += s.encode("utf-16le") + b"\0\0"
    return bytes(buf), exit_iat, rdata_rva, idt_size


def build():
    text_rva, rdata_rva, data_rva = 0x1000, 0x2000, 0x3000
    rdata, exit_iat, idir_rva, idir_size = build_rdata(rdata_rva)
    # push 0 ; call [ExitProcess] ; int3 padding
    text = b"\x6a\x00\xff\x15" + struct.pack("<I", IMAGE_BASE + exit_iat) + b"\xcc" * 8
    data = b"\0" * 0x40 + b"config-placeholder\0"

    secs = [(b".text", text_rva, text, 0x60000020), (b".rdata", rdata_rva, rdata, 0x40000040), (b".data", data_rva, data, 0xC0000040)]
    hdr_size = 0x400
    raw_ptr = hdr_size
    sec_hdrs, body = b"", b""
    for name, rva, content, flags in secs:
        raw_size = align(len(content), FILE_ALIGN)
        sec_hdrs += struct.pack("<8sIIIIIIHHI", name.ljust(8, b"\0"), len(content), rva, raw_size, raw_ptr, 0, 0, 0, 0, flags)
        body += content.ljust(raw_size, b"\0")
        raw_ptr += raw_size
    size_of_image = align(data_rva + len(data), SEC_ALIGN)

    dos = bytearray(b"MZ" + b"\0" * 62)
    struct.pack_into("<I", dos, 0x3C, 0x80)
    stub = b"\x0e\x1f\xba\x0e\x00\xb4\x09\xcd\x21\xb8\x01\x4c\xcd\x21This program cannot be run in DOS mode.\r\r\n$"
    dos = bytes(dos) + stub.ljust(0x80 - 64, b"\0")

    coff = b"PE\0\0" + struct.pack("<HHIIIHH", 0x14C, len(secs), TIMESTAMP, 0, 0, 0xE0, 0x0102)
    opt = struct.pack("<HBBIIIIIIIIIHHHHHHIIIIHHIIIIII",
                      0x10B, 14, 0, align(len(text), FILE_ALIGN), align(len(rdata) + len(data), FILE_ALIGN), 0,
                      text_rva, text_rva, rdata_rva, IMAGE_BASE, SEC_ALIGN, FILE_ALIGN,
                      6, 0, 1, 0, 6, 0, 0, size_of_image, hdr_size, 0, 3, 0x8140,
                      0x100000, 0x1000, 0x100000, 0x1000, 0, 16)
    dirs = [(0, 0)] * 16
    dirs[1] = (idir_rva, idir_size)
    dirs[12] = (rdata_rva + 20 * (len(IMPORTS) + 1) + 4 * (len(IMPORTS[b"KERNEL32.dll"]) + 1), 4 * (len(IMPORTS[b"KERNEL32.dll"]) + 1))  # IAT (first dll) - informational
    opt += b"".join(struct.pack("<II", *d) for d in dirs)
    headers = (dos + coff + opt + sec_hdrs).ljust(hdr_size, b"\0")
    overlay = b"OVERLAY-TEST-DATA:" + bytes(range(256)) * 2
    return headers + body + overlay


def main():
    out = sys.argv[1] if len(sys.argv) > 1 else "invoice_2024.exe"
    data = build()
    assert data[:2] == b"MZ" and data[0x80:0x84] == b"PE\0\0"
    with open(out, "wb") as f:
        f.write(data)
    print(f"wrote {out} ({len(data)} bytes)")


if __name__ == "__main__":
    main()