#!/usr/bin/env python3
"""
Tiny helper for static RE of GC/game471.exe:
- Find string xrefs (VA pointers to string bytes in the image)
- Find IAT callsites by scanning for "call [abs]" where abs == IAT VA
- Dump capstone disassembly around hits (no decompilation)

This is intended for format/behavior understanding (clean-room), not patching.
"""

from __future__ import annotations

import argparse
import os
import struct
from dataclasses import dataclass
from typing import Iterable, List, Optional, Tuple

import pefile  # type: ignore
from capstone import Cs, CS_ARCH_X86, CS_MODE_32  # type: ignore


@dataclass(frozen=True)
class Section:
    name: str
    rva: int
    vsize: int
    raw: int
    raw_size: int

    def contains_rva(self, rva: int) -> bool:
        return self.rva <= rva < self.rva + max(self.vsize, self.raw_size)

    def contains_raw(self, off: int) -> bool:
        return self.raw <= off < self.raw + self.raw_size


@dataclass(frozen=True)
class PeImage:
    path: str
    blob: bytes
    image_base: int
    entry_rva: int
    sections: List[Section]

    def rva_to_raw(self, rva: int) -> Optional[int]:
        for s in self.sections:
            if s.contains_rva(rva):
                return s.raw + (rva - s.rva)
        return None

    def raw_to_rva(self, off: int) -> Optional[int]:
        for s in self.sections:
            if s.contains_raw(off):
                return s.rva + (off - s.raw)
        return None

    def va_to_raw(self, va: int) -> Optional[int]:
        return self.rva_to_raw(va - self.image_base)

    def raw_to_va(self, off: int) -> Optional[int]:
        rva = self.raw_to_rva(off)
        return None if rva is None else self.image_base + rva


def load_pe(path: str) -> PeImage:
    with open(path, "rb") as f:
        blob = f.read()
    pe = pefile.PE(path, fast_load=False)
    sections: List[Section] = []
    for s in pe.sections:
        name = s.Name.rstrip(b"\x00").decode("ascii", "replace")
        sections.append(
            Section(
                name=name,
                rva=int(s.VirtualAddress),
                vsize=int(s.Misc_VirtualSize),
                raw=int(s.PointerToRawData),
                raw_size=int(s.SizeOfRawData),
            )
        )
    return PeImage(
        path=path,
        blob=blob,
        image_base=int(pe.OPTIONAL_HEADER.ImageBase),
        entry_rva=int(pe.OPTIONAL_HEADER.AddressOfEntryPoint),
        sections=sections,
    )


def find_bytes(haystack: bytes, needle: bytes) -> List[int]:
    out: List[int] = []
    pos = 0
    while True:
        i = haystack.find(needle, pos)
        if i < 0:
            return out
        out.append(i)
        pos = i + 1


def iter_xrefs_to_va(img: PeImage, target_va: int) -> Iterable[Tuple[int, int]]:
    # Returns (xref_raw_off, xref_va) where xref contains little-endian dword == target_va.
    pat = struct.pack("<I", target_va & 0xFFFFFFFF)
    for raw_off in find_bytes(img.blob, pat):
        va = img.raw_to_va(raw_off)
        if va is None:
            continue
        yield (raw_off, va)


def disasm_window(img: PeImage, va: int, before: int, after: int) -> List[str]:
    md = Cs(CS_ARCH_X86, CS_MODE_32)
    md.detail = False
    # Heuristic: try to start at a function prolog to avoid decoding garbage.
    start_va = max(img.image_base, va - before)
    end_va = va + after
    start_off = img.va_to_raw(start_va)
    end_off = img.va_to_raw(end_va)
    if start_off is None or end_off is None or end_off <= start_off:
        return [f"(cannot map VA window {hex(start_va)}..{hex(end_va)})"]

    # Scan backward for common prologs.
    scan_start = start_off
    scan_end = img.va_to_raw(va)
    if scan_end is None:
        scan_end = start_off
    scan = img.blob[scan_start:scan_end]
    prologs = [b"\x55\x8b\xec", b"\x8b\xff\x55\x8b\xec"]
    best = None
    for p in prologs:
        j = scan.rfind(p)
        if j >= 0:
            cand = scan_start + j
            if best is None or cand > best:
                best = cand
    if best is not None:
        start_off = best
        start_va2 = img.raw_to_va(start_off)
        if start_va2 is not None:
            start_va = start_va2

    code = img.blob[start_off:end_off]
    lines: List[str] = []
    for ins in md.disasm(code, start_va):
        lines.append(f"{ins.address:08x}: {ins.mnemonic:6s} {ins.op_str}")
    if not lines:
        lines.append("(no instructions decoded)")
    return lines


def find_string_va(img: PeImage, s: bytes) -> Optional[int]:
    # Find first occurrence in the file and map to VA if inside a section.
    # Some strings are stored as "...\n\0" rather than a direct "\0" terminator.
    idx = img.blob.find(s)
    if idx < 0:
        # Try common variants.
        if not s.endswith(b"\x00"):
            idx = img.blob.find(s + b"\x00")
        if idx < 0 and not s.endswith(b"\n"):
            idx = img.blob.find(s + b"\n\x00")
        if idx < 0 and not s.endswith(b"\x00"):
            idx = img.blob.find(s + b"\n")
        if idx < 0:
            return None
    rva = img.raw_to_rva(idx)
    if rva is None:
        return None
    return img.image_base + rva


def list_import_vas(path: str) -> dict[str, int]:
    pe = pefile.PE(path, fast_load=False)
    pe.parse_data_directories(
        directories=[
            pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"],
        ]
    )
    out: dict[str, int] = {}
    for entry in getattr(pe, "DIRECTORY_ENTRY_IMPORT", []):
        dll = entry.dll.decode("ascii", "ignore").lower()
        for imp in entry.imports:
            if not imp.name:
                continue
            name = imp.name.decode("ascii", "ignore")
            key = f"{dll}!{name}"
            # imp.address is the VA of the IAT slot (already ImageBase + RVA)
            out[key] = int(imp.address)
    return out


def find_iat_callsites(img: PeImage, iat_va: int) -> List[int]:
    """
    Find call dword ptr [abs32] with abs32 == iat_va.
    x86 encoding: FF 15 <imm32>
    """
    pat = b"\xFF\x15" + struct.pack("<I", iat_va & 0xFFFFFFFF)
    hits = find_bytes(img.blob, pat)
    # Return VAs of the *instruction* where it occurs.
    out: List[int] = []
    for off in hits:
        va = img.raw_to_va(off)
        if va is not None:
            out.append(va)
    return out


def find_rel32_callsites_to(img: PeImage, target_va: int) -> List[int]:
    """
    Find `E8 rel32` calls that land at target_va. Returns callsite VA.
    """
    out: List[int] = []
    blob = img.blob
    # Limit scan to .text section for speed and relevance.
    text = next((s for s in img.sections if s.name == ".text"), None)
    if not text:
        return out
    start = text.raw
    end = text.raw + text.raw_size
    i = start
    while i + 5 <= end:
        if blob[i] == 0xE8:
            rel = struct.unpack_from("<i", blob, i + 1)[0]
            callsite_va = img.raw_to_va(i)
            if callsite_va is not None:
                next_va = callsite_va + 5
                dest = (next_va + rel) & 0xFFFFFFFF
                if dest == (target_va & 0xFFFFFFFF):
                    out.append(callsite_va)
            i += 5
        else:
            i += 1
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("exe", help="Path to game471.exe")
    ap.add_argument("--string", default="data/stage/%s.dat", help="String to xref (ASCII).")
    ap.add_argument("--before", type=int, default=0x80, help="Bytes before for disasm window.")
    ap.add_argument("--after", type=int, default=0x180, help="Bytes after for disasm window.")
    ap.add_argument(
        "--imports",
        action="store_true",
        help="Also locate IAT callsites for CreateFileA/ReadFile/SetFilePointer/GetFileSize/CloseHandle.",
    )
    ap.add_argument(
        "--call-xrefs",
        default="",
        help="Find rel32 E8 callsites to this VA (hex, e.g. 0x65bb00).",
    )
    ap.add_argument(
        "--va-xrefs",
        default="",
        help="Find dword xrefs to this VA by scanning for a literal little-endian pointer (hex, e.g. 0x4ca320).",
    )
    args = ap.parse_args()

    img = load_pe(args.exe)
    print(f"exe: {img.exe if hasattr(img,'exe') else img.path}")
    print(f"image_base={hex(img.image_base)} entry_va={hex(img.image_base + img.entry_rva)}")
    print("sections:")
    for s in img.sections:
        print(f"  {s.name:8s} rva={hex(s.rva)} raw={hex(s.raw)} size={hex(s.raw_size)}")

    needle = args.string.encode("ascii")
    s_va = find_string_va(img, needle)
    if s_va is None:
        print(f"string not found: {args.string!r}")
    else:
        print(f"\nstring {args.string!r} VA={hex(s_va)}")
        xrefs = list(iter_xrefs_to_va(img, s_va))
        print(f"xrefs: {len(xrefs)}")
        for raw_off, xref_va in xrefs[:20]:
            print(f"\n== xref @ {hex(xref_va)} (raw {hex(raw_off)}) ==")
            for line in disasm_window(img, xref_va, args.before, args.after)[:80]:
                print(line)

    if args.imports:
        imports = list_import_vas(args.exe)
        want = [
            "kernel32.dll!CreateFileA",
            "kernel32.dll!ReadFile",
            "kernel32.dll!SetFilePointer",
            "kernel32.dll!GetFileSize",
            "kernel32.dll!CloseHandle",
        ]
        for k in want:
            va = imports.get(k)
            if va is None:
                print(f"\nIAT: missing {k}")
                continue
            hits = find_iat_callsites(img, va)
            print(f"\nIAT {k} slot={hex(va)} callsites={len(hits)}")
            for h in hits[:10]:
                print(f"  callsite {hex(h)}")
            if hits:
                print(f"\n== {k} first callsite disasm ==")
                for line in disasm_window(img, hits[0], args.before, args.after)[:80]:
                    print(line)

    if args.call_xrefs:
        target_va = int(args.call_xrefs, 16)
        hits = find_rel32_callsites_to(img, target_va)
        print(f"\nrel32 call xrefs to {hex(target_va)}: {len(hits)}")
        for h in hits[:20]:
            print(f"  callsite {hex(h)}")
        if hits:
            print("\n== first callsite disasm ==")
            for line in disasm_window(img, hits[0], args.before, args.after)[:80]:
                print(line)

    if args.va_xrefs:
        target_va = int(args.va_xrefs, 16)
        xrefs = list(iter_xrefs_to_va(img, target_va))
        print(f"\ndword xrefs to {hex(target_va)}: {len(xrefs)}")
        for raw_off, xref_va in xrefs[:20]:
            print(f"\n== xref @ {hex(xref_va)} (raw {hex(raw_off)}) ==")
            for line in disasm_window(img, xref_va, args.before, args.after)[:80]:
                print(line)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())