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

v9: hook CPhysicsObj::Destroy to call unparent_children(this) first.

Mechanism:
  Insert a 5-byte JMP at CPhysicsObj::Destroy entry (0x005145D0) that
  redirects to a 17-byte thunk in the .text cave at 0x00792CE0.
  The thunk:
    1. Saves this (ecx)
    2. Calls CPhysicsObj::unparent_children(this) at 0x00513FE0
    3. Restores this
    4. Re-executes the 5 displaced bytes (push ebx; push esi; mov esi,ecx; push edi)
    5. JMPs back to 0x005145D5 (continuation of Destroy)

This fixes the orphan-children leak: parents being destroyed without
calling unparent_children leave children with parent pointers to freed
memory, which causes UAF crashes in CObjectMaint cleanup chains.

Risk:
  - unparent_children is idempotent (safe to call multiple times)
  - Both patch site and thunk are in .text
  - Single chokepoint: only ~CPhysicsObj reaches Destroy
  - Stack alignment preserved (push ecx + pop ecx)
"""
import argparse
import ctypes
import ctypes.wintypes as wt
import struct
import sys


PATCH_SITE_VA  = 0x005145D0
THUNK_VA       = 0x00792CE0
UNPARENT_VA    = 0x00513FE0
RESUME_VA      = 0x005145D5

ORIG_BYTES = bytes([0x53, 0x56, 0x8B, 0xF1, 0x57])   # push ebx; push esi; mov esi,ecx; push edi
# Replacement is computed at runtime (E9 + rel32 to thunk)


def build_thunk(thunk_base: int) -> bytes:
    """17 bytes: save ecx, call unparent_children, restore, displaced, jmp back."""
    out = bytearray()
    out += bytes([0x51])                                # push ecx
    # call rel32 unparent_children
    rel_call = UNPARENT_VA - (thunk_base + len(out) + 5)
    out += bytes([0xE8]) + struct.pack("<i", rel_call)
    out += bytes([0x59])                                # pop ecx
    # Displaced prologue bytes (5 bytes)
    out += ORIG_BYTES                                    # push ebx; push esi; mov esi,ecx; push edi
    # jmp rel32 back to RESUME_VA
    rel_jmp = RESUME_VA - (thunk_base + len(out) + 5)
    out += bytes([0xE9]) + struct.pack("<i", rel_jmp)
    return bytes(out)


PROCESS_VM_READ           = 0x0010
PROCESS_VM_WRITE          = 0x0020
PROCESS_VM_OPERATION      = 0x0008
PROCESS_QUERY_INFORMATION = 0x0400
PAGE_EXECUTE_READWRITE    = 0x40


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


def read_bytes(h, addr, n):
    buf = (ctypes.c_ubyte * n)()
    sz = ctypes.c_size_t(0)
    if not ReadProcessMemory(h, addr, buf, n, ctypes.byref(sz)):
        raise OSError(f"read 0x{addr:08x} err={ctypes.get_last_error()}")
    return bytes(buf[:sz.value])


def write_bytes(h, addr, data):
    old_prot = wt.DWORD(0)
    if not VirtualProtectEx(h, addr, len(data), PAGE_EXECUTE_READWRITE, ctypes.byref(old_prot)):
        raise OSError(f"VirtualProtectEx 0x{addr:08x} err={ctypes.get_last_error()}")
    sz = ctypes.c_size_t(0)
    ok = WriteProcessMemory(h, addr, data, len(data), ctypes.byref(sz))
    err = ctypes.get_last_error() if not ok else 0
    restored = wt.DWORD(0)
    VirtualProtectEx(h, addr, len(data), old_prot.value, ctypes.byref(restored))
    if not ok:
        raise OSError(f"write 0x{addr:08x} err={err}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("--revert", action="store_true")
    args = ap.parse_args()

    h = OpenProcess(
        PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_QUERY_INFORMATION,
        False, args.pid,
    )
    if not h:
        print(f"OpenProcess({args.pid}) err={ctypes.get_last_error()}"); sys.exit(2)

    cur = read_bytes(h, PATCH_SITE_VA, 5)
    print(f"PID {args.pid}")
    print(f"  patch site @ 0x{PATCH_SITE_VA:08x} current: {cur.hex()}")

    if args.revert:
        if cur == ORIG_BYTES:
            print(f"    already original")
        else:
            # Restore the original 5 bytes
            write_bytes(h, PATCH_SITE_VA, ORIG_BYTES)
            after = read_bytes(h, PATCH_SITE_VA, 5)
            print(f"    reverted; bytes now: {after.hex()}")
        CloseHandle(h); return

    if cur != ORIG_BYTES:
        print(f"    UNEXPECTED — bytes don't match expected original {ORIG_BYTES.hex()}")
        CloseHandle(h); sys.exit(3)

    # Verify cave is clean (zero-filled or at least not in-use)
    cave_cur = read_bytes(h, THUNK_VA, 32)
    print(f"  cave @ 0x{THUNK_VA:08x} current (32B): {cave_cur.hex()}")
    if any(b != 0 and b != 0xCC for b in cave_cur):
        # Warn but don't refuse — caves are sometimes filled with INT3 (0xCC) padding
        print(f"    WARNING — cave is not all zeros/0xCC. Inspect before proceeding.")
        # For safety, still refuse if first byte is something live
        if cave_cur[0] != 0 and cave_cur[0] != 0xCC:
            print(f"    REFUSING — cave appears occupied")
            CloseHandle(h); sys.exit(4)

    # Build and write thunk
    thunk = build_thunk(THUNK_VA)
    assert len(thunk) == 17, f"thunk size {len(thunk)}"
    print(f"  writing thunk ({len(thunk)} bytes) to 0x{THUNK_VA:08x}: {thunk.hex()}")
    write_bytes(h, THUNK_VA, thunk)
    # Verify thunk
    after_thunk = read_bytes(h, THUNK_VA, len(thunk))
    if after_thunk != thunk:
        print(f"    THUNK MISMATCH after write: {after_thunk.hex()}")
        CloseHandle(h); sys.exit(5)
    print(f"    thunk verified")

    # Write JMP at patch site
    rel = THUNK_VA - (PATCH_SITE_VA + 5)
    jmp_bytes = bytes([0xE9]) + struct.pack("<i", rel)
    print(f"  writing JMP at 0x{PATCH_SITE_VA:08x}: {jmp_bytes.hex()}")
    write_bytes(h, PATCH_SITE_VA, jmp_bytes)
    after = read_bytes(h, PATCH_SITE_VA, 5)
    print(f"    patch site now: {after.hex()}")
    if after != jmp_bytes:
        print(f"    PATCH MISMATCH"); sys.exit(6)

    print(f"  OK")
    CloseHandle(h)


if __name__ == "__main__":
    main()