#!/usr/bin/env python3
"""Summarize PE imports and simple IAT references for Hwanse2.exe."""
from __future__ import annotations

import argparse
import json
import struct
from pathlib import Path


IMAGE_DIRECTORY_ENTRY_IMPORT = 1
IMAGE_ORDINAL_FLAG32 = 0x80000000


def c_string(data: bytes, offset: int) -> str:
    end = data.find(b"\x00", offset)
    if end < 0:
        end = len(data)
    return data[offset:end].decode("ascii", errors="replace")


def pe_headers(data: bytes) -> tuple[int, int, list[dict]]:
    if data[:2] != b"MZ":
        raise ValueError("not an MZ executable")
    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if data[pe_offset:pe_offset + 4] != b"PE\x00\x00":
        raise ValueError("not a PE executable")

    file_header = pe_offset + 4
    section_count = struct.unpack_from("<H", data, file_header + 2)[0]
    optional_size = struct.unpack_from("<H", data, file_header + 16)[0]
    optional = file_header + 20
    magic = struct.unpack_from("<H", data, optional)[0]
    if magic != 0x10B:
        raise ValueError("only PE32 is supported")

    image_base = struct.unpack_from("<I", data, optional + 28)[0]
    import_rva, import_size = struct.unpack_from("<II", data, optional + 96 + IMAGE_DIRECTORY_ENTRY_IMPORT * 8)
    sections = []
    section_table = optional + optional_size
    for index in range(section_count):
        offset = section_table + index * 40
        name = data[offset:offset + 8].split(b"\x00", 1)[0].decode("ascii", errors="replace")
        virtual_size, virtual_address, raw_size, raw_ptr = struct.unpack_from("<IIII", data, offset + 8)
        sections.append(
            {
                "name": name,
                "virtualAddress": virtual_address,
                "virtualSize": virtual_size,
                "rawSize": raw_size,
                "rawPointer": raw_ptr,
            }
        )
    return image_base, import_rva, sections


def rva_to_offset(rva: int, sections: list[dict]) -> int:
    for section in sections:
        start = section["virtualAddress"]
        size = max(section["virtualSize"], section["rawSize"])
        if start <= rva < start + size:
            return section["rawPointer"] + (rva - start)
    raise ValueError(f"RVA 0x{rva:08x} is outside known sections")


def section_bytes(data: bytes, section: dict) -> bytes:
    start = section["rawPointer"]
    return data[start:start + section["rawSize"]]


def find_iat_refs(data: bytes, sections: list[dict], image_base: int, iat_rva: int) -> list[dict]:
    needle = struct.pack("<I", image_base + iat_rva)
    refs = []
    for section in sections:
        if section["name"] not in {".text", ".rdata", ".data"}:
            continue
        body = section_bytes(data, section)
        index = body.find(needle)
        while index >= 0:
            ref_index = index - 2 if index >= 2 and body[index - 2:index] == b"\xff\x25" else index
            refs.append(
                {
                    "section": section["name"],
                    "rva": section["virtualAddress"] + ref_index,
                    "va": image_base + section["virtualAddress"] + ref_index,
                }
            )
            index = body.find(needle, index + 1)
    return refs


def find_relative_calls(data: bytes, sections: list[dict], image_base: int, target_va: int) -> list[dict]:
    refs = []
    for section in sections:
        if section["name"] != ".text":
            continue
        body = section_bytes(data, section)
        for index in range(0, max(0, len(body) - 4)):
            if body[index] != 0xE8:
                continue
            rel = struct.unpack_from("<i", body, index + 1)[0]
            source_va = image_base + section["virtualAddress"] + index
            if source_va + 5 + rel == target_va:
                refs.append(
                    {
                        "section": section["name"],
                        "rva": section["virtualAddress"] + index,
                        "va": source_va,
                    }
                )
    return refs


def read_import_name(data: bytes, sections: list[dict], thunk_value: int) -> dict:
    if thunk_value & IMAGE_ORDINAL_FLAG32:
        return {"ordinal": thunk_value & 0xFFFF}
    name_offset = rva_to_offset(thunk_value, sections)
    hint = struct.unpack_from("<H", data, name_offset)[0]
    return {"hint": hint, "name": c_string(data, name_offset + 2)}


def parse_imports(path: Path) -> dict:
    data = path.read_bytes()
    image_base, import_rva, sections = pe_headers(data)
    import_offset = rva_to_offset(import_rva, sections)
    dlls = []

    descriptor_offset = import_offset
    while True:
        original_first_thunk, _time, _forwarder, name_rva, first_thunk = struct.unpack_from(
            "<IIIII",
            data,
            descriptor_offset,
        )
        if not any([original_first_thunk, name_rva, first_thunk]):
            break
        dll_name = c_string(data, rva_to_offset(name_rva, sections))
        thunk_rva = original_first_thunk or first_thunk
        thunk_offset = rva_to_offset(thunk_rva, sections)
        functions = []
        index = 0
        while True:
            thunk_value = struct.unpack_from("<I", data, thunk_offset + index * 4)[0]
            if thunk_value == 0:
                break
            imported = read_import_name(data, sections, thunk_value)
            iat_rva = first_thunk + index * 4
            imported["iatRva"] = iat_rva
            imported["iatVa"] = image_base + iat_rva
            imported["refs"] = find_iat_refs(data, sections, image_base, iat_rva)
            imported["callRefs"] = [
                call_ref
                for ref in imported["refs"]
                for call_ref in find_relative_calls(data, sections, image_base, ref["va"])
            ]
            functions.append(imported)
            index += 1
        dlls.append({"dll": dll_name, "functions": functions})
        descriptor_offset += 20

    return {
        "exe": str(path),
        "imageBase": image_base,
        "imports": dlls,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# EXE Imports",
        "",
        f"Executable: `{summary['exe']}`",
        f"Image base: `0x{summary['imageBase']:08x}`",
        "",
        "## Imported DLLs",
        "",
        "| DLL | functions | notable imports |",
        "| --- | ---: | --- |",
    ]
    for dll in summary["imports"]:
        names = [item.get("name", f"ordinal:{item.get('ordinal')}") for item in dll["functions"]]
        notable = ", ".join(names[:12])
        if len(names) > 12:
            notable += f", ... (+{len(names) - 12})"
        lines.append(f"| {dll['dll']} | {len(names)} | {notable} |")

    lines.extend(["", "## DirectX Imports", ""])
    for dll in summary["imports"]:
        if dll["dll"].upper() not in {"DDRAW.DLL", "DINPUT.DLL", "DSOUND.DLL"}:
            continue
        lines.extend([f"### {dll['dll']}", "", "| import | IAT VA | thunk refs | call refs |", "| --- | --- | --- | --- |"])
        for item in dll["functions"]:
            refs = ", ".join(f"{ref['section']}:0x{ref['va']:08x}" for ref in item["refs"][:8])
            if len(item["refs"]) > 8:
                refs += f", ... (+{len(item['refs']) - 8})"
            call_refs = ", ".join(f"{ref['section']}:0x{ref['va']:08x}" for ref in item["callRefs"][:12])
            if len(item["callRefs"]) > 12:
                call_refs += f", ... (+{len(item['callRefs']) - 12})"
            lines.append(
                f"| {item.get('name', 'ordinal:' + str(item.get('ordinal')))} | "
                f"`0x{item['iatVa']:08x}` | {refs or '-'} | {call_refs or '-'} |"
            )
        lines.append("")

    return "\n".join(lines).rstrip() + "\n"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=Path("Hwanse2.exe"))
    parser.add_argument("--json-out", type=Path, default=Path("out/exe_imports.json"))
    parser.add_argument("--md-out", "--out", dest="md_out", type=Path, default=None)
    args = parser.parse_args()

    summary = parse_imports(args.exe)
    args.json_out.parent.mkdir(parents=True, exist_ok=True)
    args.json_out.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if args.md_out:
        args.md_out.write_text(markdown(summary), encoding="utf-8")
    print(f"wrote EXE import summary -> {args.json_out}")


if __name__ == "__main__":
    main()