"""patch_v14_position_alloc_trace.py <pid> [--revert]

DIAGNOSTIC ONLY — does not modify program behavior.

Goal: find the call sites that produce leaked Position instances.

Approach: Install a counter-increment thunk at the start of each
Position::Position ctor variant. The thunk:
  - increments a per-ctor 32-bit counter in a runtime-allocated
    counter block
  - then falls through to the original prologue and the rest of
    the ctor

Counter block is queried by `--read`. No game state is modified.

CTORS (EoR addresses, derived from 2013):
  - Position::Position()             — 2013 0x00424ab0 → EoR TBD
  - Position::Position(uint, Frame*) — 2013 0x00452780 → EoR TBD
  - Position::Position(Position*)    — 2013 0x004529a0 → EoR TBD

The EoR offsets are NOT yet known precisely — this patcher is a
SKELETON. Before running, the analyst must:
  1. Find Position::Position in EoR via reference from
     `vtable = 0x00797910` writes (cdb: `s -d acclient_base
     L?<size> 00797910`)
  2. Fill in CTOR_ADDRS below with the three EoR addresses.
  3. Verify each address starts with `mov eax, [esp+4]` or similar
     stack-arg setup; check that 5 bytes of prologue fits a relative
     `call <thunk>` (e.g. `e8 XX XX XX XX`).

After verifying live counts are non-zero and growing, leave the
patcher in place for 1 hour and dump counters via `--read`. The
ratio of count_a / count_b / count_c identifies the dominant call
path.

SAFETY: the thunks preserve all registers (push eax/pushfd) before
incrementing, restore after. They do NOT change any game state.
Even if the analyst applies the wrong ctor address, the worst case
is a stale call (caller crashes immediately, easy to diagnose) — no
delayed-AV class of failure.

DO NOT apply this patcher in production. It is for read-only
behavioral diagnosis only. Apply on ONE non-essential PID (e.g.
Time, the idle character) and observe for 30 minutes before
inferring rate.
"""
import argparse
import ctypes
import ctypes.wintypes as wt
import struct
import sys

# --- CONFIG (FILL IN BEFORE USE) -----
# These are placeholders — verify against EoR binary before applying.
# In 2013 these were 0x00424ab0, 0x00452780, 0x004529a0. The EoR
# offset relative to 2013 varies per class; do NOT assume +0x1000.
CTOR_ADDRS = {
    "default":  0x00000000,   # FILL IN — 2013 0x00424ab0
    "uint_fp":  0x00000000,   # FILL IN — 2013 0x00452780
    "copy":     0x00000000,   # FILL IN — 2013 0x004529a0
}

VERIFY_VT_ADDR = 0x00797910   # Position vtable; should match
                              # `mov dword ptr [ecx], 00797910h`
                              # in each ctor prologue.

# --- Win32 surface --------------------
PROCESS_ALL_ACCESS = 0x1F0FFF
MEM_COMMIT = 0x1000
MEM_RESERVE = 0x2000
PAGE_EXECUTE_READWRITE = 0x40
PAGE_READWRITE = 0x04

k = ctypes.windll.kernel32
k.OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]
k.OpenProcess.restype = wt.HANDLE
k.CloseHandle.argtypes = [wt.HANDLE]
k.CloseHandle.restype = wt.BOOL
k.ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID,
                                ctypes.c_size_t,
                                ctypes.POINTER(ctypes.c_size_t)]
k.ReadProcessMemory.restype = wt.BOOL
k.WriteProcessMemory.argtypes = [wt.HANDLE, wt.LPVOID, wt.LPCVOID,
                                 ctypes.c_size_t,
                                 ctypes.POINTER(ctypes.c_size_t)]
k.WriteProcessMemory.restype = wt.BOOL
k.VirtualAllocEx.argtypes = [wt.HANDLE, wt.LPVOID, ctypes.c_size_t,
                             wt.DWORD, wt.DWORD]
k.VirtualAllocEx.restype = wt.LPVOID
k.VirtualFreeEx.argtypes = [wt.HANDLE, wt.LPVOID, ctypes.c_size_t,
                            wt.DWORD]
k.VirtualFreeEx.restype = wt.BOOL


def rpm(h, addr, n):
    buf = (ctypes.c_ubyte * n)()
    sz = ctypes.c_size_t(0)
    if not k.ReadProcessMemory(h, addr, buf, n, ctypes.byref(sz)):
        return None
    return bytes(buf[:sz.value])


def wpm(h, addr, data):
    buf = (ctypes.c_ubyte * len(data))(*data)
    sz = ctypes.c_size_t(0)
    return bool(k.WriteProcessMemory(h, addr, buf, len(data),
                                     ctypes.byref(sz)))


def verify_ctor_at(h, addr):
    """Confirm the ctor at `addr` is a Position ctor by looking for
    a write of vt 0x00797910 in the first 32 bytes."""
    code = rpm(h, addr, 32)
    if not code:
        return False
    target = struct.pack("<I", VERIFY_VT_ADDR)
    return target in code


def apply(pid):
    if 0 in CTOR_ADDRS.values():
        print("ERROR: CTOR_ADDRS not filled in. Edit the constants at "
              "the top of this script first.")
        return 1
    h = k.OpenProcess(PROCESS_ALL_ACCESS, False, pid)
    if not h:
        print(f"OpenProcess failed (err={ctypes.get_last_error()})")
        return 1
    try:
        # Verify each ctor address writes the Position vt
        for label, addr in CTOR_ADDRS.items():
            if not verify_ctor_at(h, addr):
                print(f"ERROR: addr 0x{addr:08x} ({label}) does not "
                      f"reference Position vt 0x{VERIFY_VT_ADDR:08x}. "
                      "Refusing to patch.")
                return 2

        # Allocate counter block + thunk region
        block = k.VirtualAllocEx(h, None, 0x1000,
                                 MEM_COMMIT | MEM_RESERVE,
                                 PAGE_EXECUTE_READWRITE)
        if not block:
            print("VirtualAllocEx failed")
            return 3
        block = int(block)
        print(f"counter+thunk block @ 0x{block:08x}")

        # Layout:
        #   block + 0x000   counter A (default ctor)
        #   block + 0x004   counter B (uint+frame ctor)
        #   block + 0x008   counter C (copy ctor)
        #   block + 0x100   thunk A
        #   block + 0x120   thunk B
        #   block + 0x140   thunk C
        # Each thunk:
        #   pushfd; push eax
        #   mov eax, &counter
        #   lock inc [eax]
        #   pop eax; popfd
        #   <copy of overwritten prologue bytes>
        #   jmp <ctor + 5>
        labels = list(CTOR_ADDRS.keys())
        for i, label in enumerate(labels):
            ctor_addr = CTOR_ADDRS[label]
            counter_addr = block + i * 4
            thunk_addr = block + 0x100 + i * 0x20
            # Read 5 prologue bytes (the bytes we'll relocate)
            prologue = rpm(h, ctor_addr, 5)
            if not prologue or len(prologue) != 5:
                print(f"failed to read prologue at 0x{ctor_addr:08x}")
                return 4
            # Build thunk
            thunk = bytearray()
            thunk += b"\x9c"                          # pushfd
            thunk += b"\x50"                          # push eax
            thunk += b"\xb8" + struct.pack("<I", counter_addr)  # mov eax, &counter
            thunk += b"\xf0\xff\x00"                  # lock inc dword [eax]
            thunk += b"\x58"                          # pop eax
            thunk += b"\x9d"                          # popfd
            thunk += prologue                         # original 5 bytes
            # Jump back to ctor+5
            return_to = ctor_addr + 5
            rel = return_to - (thunk_addr + len(thunk) + 5)
            thunk += b"\xe9" + struct.pack("<i", rel)  # jmp rel32
            wpm(h, thunk_addr, bytes(thunk))
            # Write `e9 rel32` (call would be cleaner; use jmp since
            # we replicate prologue inside thunk)
            patch_rel = thunk_addr - (ctor_addr + 5)
            patch = b"\xe9" + struct.pack("<i", patch_rel)
            wpm(h, ctor_addr, patch)
            print(f"  patched {label} at 0x{ctor_addr:08x} → "
                  f"thunk 0x{thunk_addr:08x}, counter @ 0x{counter_addr:08x}")
        print(f"\nCounters at 0x{block:08x}. Read with `--read 0x{block:08x}`.")
        return 0
    finally:
        k.CloseHandle(h)


def read_counters(pid, block):
    h = k.OpenProcess(PROCESS_ALL_ACCESS, False, pid)
    if not h:
        print("OpenProcess failed")
        return 1
    try:
        data = rpm(h, block, 16)
        if not data:
            print("read failed")
            return 1
        a, b, c, _ = struct.unpack("<IIII", data)
        print(f"  default ctor:    {a:>10}")
        print(f"  uint+frame ctor: {b:>10}")
        print(f"  copy ctor:       {c:>10}")
        print(f"  total:           {a+b+c:>10}")
        return 0
    finally:
        k.CloseHandle(h)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("--read", type=lambda x: int(x, 0),
                    help="read counters at this address")
    ap.add_argument("--revert", action="store_true",
                    help="NOT IMPLEMENTED - this is a one-shot "
                         "diagnostic; restart client to revert")
    args = ap.parse_args()
    if args.read:
        return read_counters(args.pid, args.read)
    if args.revert:
        print("revert not implemented — kill client to undo")
        return 1
    return apply(args.pid)


if __name__ == "__main__":
    sys.exit(main())