"""add_import.py <input.exe> <output.exe> <dll_name>

Patch a PE EXE's import table to add a new DLL import.

The OS loader will pull <dll_name> into the process before the EXE's
entry point runs — exactly what we want for leakfix.dll.

Mechanism:
  1. Read the PE file.
  2. Add a new section called ".limport" at the end with:
        - new IMAGE_IMPORT_DESCRIPTOR array (existing entries + ours + null)
        - ILT (Import Lookup Table) and IAT for our DLL — both pointing
          at a single hint/name "LeakfixStub" (any name; doesn't have to
          exist since loading the DLL is enough to trigger its DllMain).
        - The DLL name string.
        - Hint/name table for our exported function.
  3. Update OptionalHeader.DataDirectory[1] (IMPORT) to point at our
     new array, with the size covering all entries.
  4. Write the new file.

We must pick an export name that exists in leakfix.dll for the loader
to resolve at load time, OR we can use ordinal #1 if we export by
ordinal. Simplest: have leakfix.dll export a stub function named
"leakfix_init" via __declspec(dllexport), and reference that here.
"""
import struct, sys, os

PE_MACHINE_I386 = 0x014c

def u8(b, o): return b[o]
def u16(b, o): return struct.unpack_from("<H", b, o)[0]
def u32(b, o): return struct.unpack_from("<I", b, o)[0]

def main():
    if len(sys.argv) != 4:
        print(__doc__); sys.exit(1)
    inp, outp, dll_name = sys.argv[1], sys.argv[2], sys.argv[3]

    with open(inp, "rb") as f:
        data = bytearray(f.read())

    # 1. Locate headers
    if data[:2] != b"MZ":
        print("not a PE file"); sys.exit(2)
    pe_off = u32(data, 0x3c)
    if data[pe_off:pe_off+4] != b"PE\0\0":
        print("PE signature not found"); sys.exit(2)
    machine = u16(data, pe_off + 4)
    if machine != PE_MACHINE_I386:
        print(f"unexpected machine 0x{machine:04x} (want 0x14c = i386)"); sys.exit(2)
    num_sections = u16(data, pe_off + 6)
    size_of_optional = u16(data, pe_off + 20)
    optional_off = pe_off + 24
    section_table_off = optional_off + size_of_optional

    # PE32 (not PE32+); confirm magic 0x10b
    if u16(data, optional_off) != 0x010b:
        print("not PE32 (32-bit) optional header magic"); sys.exit(2)

    image_base   = u32(data, optional_off + 28)
    section_align = u32(data, optional_off + 32)
    file_align    = u32(data, optional_off + 36)
    size_of_image = u32(data, optional_off + 56)
    size_of_headers = u32(data, optional_off + 60)
    data_dir_off = optional_off + 96  # for PE32

    # Existing IMPORT directory
    imp_rva = u32(data, data_dir_off + 1*8)
    imp_size = u32(data, data_dir_off + 1*8 + 4)
    print(f"PE32 image_base=0x{image_base:08x}, sectionAlign=0x{section_align:x}, fileAlign=0x{file_align:x}")
    print(f"existing IMPORT dir: rva=0x{imp_rva:08x} size={imp_size}")

    # 2. Read all sections
    sections = []
    for i in range(num_sections):
        sh = section_table_off + i * 40
        name = bytes(data[sh:sh+8]).rstrip(b"\0").decode("ascii", "replace")
        vsize = u32(data, sh+8)
        vaddr = u32(data, sh+12)
        rsize = u32(data, sh+16)
        roff  = u32(data, sh+20)
        chars = u32(data, sh+36)
        sections.append({"name":name, "vsize":vsize, "vaddr":vaddr, "rsize":rsize, "roff":roff, "chars":chars, "sh_off":sh})

    # find rva-to-file mapping for IMPORT
    def rva_to_off(rva):
        for s in sections:
            if s["vaddr"] <= rva < s["vaddr"] + max(s["vsize"], s["rsize"]):
                return s["roff"] + (rva - s["vaddr"])
        return None

    imp_off = rva_to_off(imp_rva)
    if imp_off is None: print("can't map import RVA"); sys.exit(2)

    # 3. Count existing import descriptors (each is 20 bytes; terminated by zero descriptor)
    DESC_SZ = 20
    existing_descs = bytearray()
    pos = imp_off
    while True:
        d = bytes(data[pos:pos+DESC_SZ])
        if d == b"\0"*DESC_SZ: break
        existing_descs += d
        pos += DESC_SZ
    n_existing = len(existing_descs) // DESC_SZ
    print(f"existing imports: {n_existing}")

    # 4. Build new section
    # Section layout (offsets within section start):
    #   0x00  new descriptor array: existing descs + our desc + zero terminator
    #   then ILT: one DWORD pointing at name-table; one DWORD zero (terminator)
    #   then IAT: same shape
    #   then name-table: hint(2) + "leakfix_init\0"
    #   then dll-name: "leakfix.dll\0"
    new_section_align = section_align
    new_section = bytearray()

    # We don't know final RVAs yet; lay out, then patch RVAs at the end.
    n_descs = n_existing + 2          # existing + ours + terminator
    desc_table_size = n_descs * DESC_SZ

    ilt_off = desc_table_size                  # 2 DWORDs (1 hint+name RVA, 1 terminator)
    iat_off = ilt_off + 8
    name_table_off = iat_off + 8
    func_name = b"leakfix_init\0"
    # IMAGE_IMPORT_BY_NAME = WORD hint + name
    name_entry = b"\x00\x00" + func_name
    if len(name_entry) & 1: name_entry += b"\0"
    dll_name_off = name_table_off + len(name_entry)
    dll_name_b = dll_name.encode("ascii") + b"\0"
    if len(dll_name_b) & 1: dll_name_b += b"\0"

    total_data_size = dll_name_off + len(dll_name_b)

    # Round section size up to fileAlign for raw, sectionAlign for virtual
    def round_up(v, a): return (v + a - 1) & ~(a - 1)
    raw_size = round_up(total_data_size, file_align)
    virt_size = round_up(total_data_size, section_align)

    # Determine new section's RVA: at end of image
    last_vend = max((s["vaddr"] + round_up(max(s["vsize"], s["rsize"]), section_align)) for s in sections)
    new_vaddr = round_up(last_vend, section_align)
    new_roff  = len(data)                       # append to end of file
    new_roff  = round_up(new_roff, file_align)
    # Pad file up to new_roff
    if new_roff > len(data):
        data += b"\0" * (new_roff - len(data))

    # Now we know RVAs. Build section bytes.
    sec = bytearray(raw_size)

    # Copy existing descriptors verbatim, then append our descriptor, then zero
    sec[0:len(existing_descs)] = existing_descs
    our_desc_off = len(existing_descs)
    # Our descriptor: ILT_RVA, TimeStamp, ForwarderChain, Name_RVA, IAT_RVA
    our_ilt_rva  = new_vaddr + ilt_off
    our_iat_rva  = new_vaddr + iat_off
    our_name_rva = new_vaddr + dll_name_off
    name_entry_rva = new_vaddr + name_table_off
    struct.pack_into("<IIIII", sec, our_desc_off,
                     our_ilt_rva, 0, 0, our_name_rva, our_iat_rva)
    # Zero terminator after our descriptor — sec is already zeroed

    # ILT/IAT entries (both point at the hint/name)
    struct.pack_into("<II", sec, ilt_off, name_entry_rva, 0)
    struct.pack_into("<II", sec, iat_off, name_entry_rva, 0)

    # Name table
    sec[name_table_off:name_table_off + len(name_entry)] = name_entry
    # DLL name
    sec[dll_name_off:dll_name_off + len(dll_name_b)] = dll_name_b

    # Append section bytes to file
    data += sec

    # 5. Update section table
    if num_sections + 1 > (size_of_headers - section_table_off + pe_off) // 40:
        # Not enough room in headers for another section entry. Bail.
        print("ERROR: no room in PE headers for an additional section entry"); sys.exit(3)

    new_sh = section_table_off + num_sections * 40
    name_bytes = b".limport"[:8].ljust(8, b"\0")
    data[new_sh:new_sh+8] = name_bytes
    struct.pack_into("<I", data, new_sh + 8,  total_data_size)         # VirtualSize
    struct.pack_into("<I", data, new_sh + 12, new_vaddr)               # VirtualAddress
    struct.pack_into("<I", data, new_sh + 16, raw_size)                # SizeOfRawData
    struct.pack_into("<I", data, new_sh + 20, new_roff)                # PointerToRawData
    struct.pack_into("<I", data, new_sh + 24, 0)                       # PointerToRelocations
    struct.pack_into("<I", data, new_sh + 28, 0)                       # PointerToLinenumbers
    struct.pack_into("<H", data, new_sh + 32, 0)                       # NumberOfRelocations
    struct.pack_into("<H", data, new_sh + 34, 0)                       # NumberOfLinenumbers
    # Characteristics: CODE? DATA. Use 0x40000040 = INITIALIZED_DATA | READ
    # We need WRITE on the IAT but for simple loaders read-only is fine
    # because the loader rewrites IAT to actual addresses (writable while loading).
    # Use IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE | INITIALIZED_DATA = 0xC0000040
    struct.pack_into("<I", data, new_sh + 36, 0xC0000040)
    # Bump section count
    struct.pack_into("<H", data, pe_off + 6, num_sections + 1)

    # 6. Update OptionalHeader: SizeOfImage, IMPORT data directory
    new_size_of_image = new_vaddr + virt_size
    struct.pack_into("<I", data, optional_off + 56, new_size_of_image)

    new_imp_rva = new_vaddr + 0    # descriptor table at start of our section
    new_imp_size = (n_existing + 1) * DESC_SZ  # not including null terminator per MS spec... but include for safety
    struct.pack_into("<II", data, data_dir_off + 1*8, new_imp_rva, new_imp_size + DESC_SZ)

    # IAT data directory (index 12) might also need updating — point at our IAT.
    # For loaders, IMPORT is what matters; IAT directory is optional. Leave alone.

    with open(outp, "wb") as f:
        f.write(data)
    print(f"wrote {outp} ({len(data)} bytes)")
    print(f"  new section @ rva 0x{new_vaddr:08x} (file 0x{new_roff:x}), size {raw_size}")
    print(f"  new IMPORT dir @ rva 0x{new_imp_rva:08x}, descriptors: {n_existing} existing + 1 ours")
    fn = func_name.rstrip(b"\0").decode()
    print(f"  added import: {dll_name} (resolves '{fn}')")

if __name__ == "__main__":
    main()