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

v12: NULL/range validator for the unpacker at 0x00526a50.

This function (presumably AnimSequenceNode UnPack overload, or similar
4-dword stream reader) is called via vtable slot at 0x007c92c8. It
dereferences `*(ctx)` without validating the cursor pointer. Two
confirmed crashes here (Shadow 20:28, Frank 23:22), both with bad
cursor pointers from leak-corrupted contexts.

Mechanism (2 sites):
  1. Write 29-byte validator at 0x00526a45 (uses 11-byte NOP pad +
     overwrites first 18 bytes of original function).
     Validator: size check + EAX load + cursor non-NULL + cursor
     >= 0x00400000 → success path jumps to 0x00526a62 (body unchanged);
     failure path returns 0 (same as original size-check-fail).
  2. Redirect vtable slot at 0x007c92c8 from 0x00526a50 → 0x00526a45.

After patch, all dispatched calls hit the validator first.
The body from 0x00526a62 onward is untouched.

Risk:
  - If a direct caller of 0x00526a50 exists (not via vtable), it'd
    land mid-validator at MOV EDX,[EAX] with EAX uninitialized.
    Ghidra found only the vtable xref, so unlikely.
  - Writing to a vtable in process memory is fine (we VirtualProtectEx).
  - 11 NOPs + 18 original-entry bytes = 29 bytes total replacement.
    No overlap into the body at 0x00526a62.
"""
import argparse, ctypes, ctypes.wintypes as wt, struct, sys


VALIDATOR_VA   = 0x00526a45
DISPATCH_VA    = 0x007c92c8
OLD_FUNC_VA    = 0x00526a50

VALIDATOR_BYTES = bytes([
    # 0x00526a45  CMP DWORD [ESP+8], 0x10
    0x83, 0x7C, 0x24, 0x08, 0x10,
    # 0x00526a4a  JB +0x11 → 0x00526a5d (fail)
    0x72, 0x11,
    # 0x00526a4c  MOV EAX, [ESP+4]  (EAX = ctx)
    0x8B, 0x44, 0x24, 0x04,
    # 0x00526a50  MOV EDX, [EAX]    (EDX = cursor)
    0x8B, 0x10,
    # 0x00526a52  CMP EDX, 0x00400000  (cursor >= image base?)
    0x81, 0xFA, 0x00, 0x00, 0x40, 0x00,
    # 0x00526a58  JB +0x03 → 0x00526a5d (fail)
    0x72, 0x03,
    # 0x00526a5a  JMP +0x06 → 0x00526a62 (body)
    0xEB, 0x06,
    # 0x00526a5c  NOP (filler)
    0x90,
    # 0x00526a5d  XOR EAX, EAX
    0x33, 0xC0,
    # 0x00526a5f  RET 8
    0xC2, 0x08, 0x00,
])
assert len(VALIDATOR_BYTES) == 29

# Original bytes at validator site: 11 NOPs + 18 bytes of original function entry
ORIG_VALIDATOR = bytes([0x90] * 11) + bytes([
    0x83, 0x7C, 0x24, 0x08, 0x10,           # CMP [ESP+8], 0x10
    0x73, 0x05,                              # JAE +5
    0x33, 0xC0,                              # XOR EAX, EAX
    0xC2, 0x08, 0x00,                        # RET 8
    0x8B, 0x44, 0x24, 0x04,                  # MOV EAX, [ESP+4]
    0x8B, 0x10                               # MOV EDX, [EAX]
])
assert len(ORIG_VALIDATOR) == 29

OLD_DISPATCH = struct.pack("<I", OLD_FUNC_VA)
NEW_DISPATCH = struct.pack("<I", VALIDATOR_VA)


PROCESS_VM_READ = 0x10
PROCESS_VM_WRITE = 0x20
PROCESS_VM_OPERATION = 0x8
PROCESS_QUERY_INFORMATION = 0x400
PAGE_EXECUTE_READWRITE = 0x40
PAGE_READWRITE = 0x04


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, prot=PAGE_EXECUTE_READWRITE):
    old = wt.DWORD(0)
    if not VirtualProtectEx(h, addr, len(data), prot, ctypes.byref(old)):
        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.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)

    print(f"PID {args.pid}")

    # Read both sites
    cur_validator = read_bytes(h, VALIDATOR_VA, 29)
    cur_dispatch = read_bytes(h, DISPATCH_VA, 4)
    print(f"  validator @ 0x{VALIDATOR_VA:08x}: cur {cur_validator.hex()}")
    print(f"  dispatch  @ 0x{DISPATCH_VA:08x}: cur {cur_dispatch.hex()}")

    if args.revert:
        if cur_dispatch == OLD_DISPATCH:
            print(f"  already original (dispatch)")
        elif cur_dispatch == NEW_DISPATCH:
            # 1. First restore dispatch (callers go back to original function)
            write_bytes(h, DISPATCH_VA, OLD_DISPATCH)
            # 2. Then restore validator bytes
            write_bytes(h, VALIDATOR_VA, ORIG_VALIDATOR)
            print(f"  reverted both sites")
        else:
            print(f"  UNEXPECTED dispatch — refusing"); sys.exit(3)
        CloseHandle(h); return

    # Apply: validator first, then dispatch
    if cur_dispatch == NEW_DISPATCH:
        print(f"  already patched")
        CloseHandle(h); return
    if cur_dispatch != OLD_DISPATCH:
        print(f"  UNEXPECTED dispatch. Expected {OLD_DISPATCH.hex()}")
        CloseHandle(h); sys.exit(4)
    if cur_validator != ORIG_VALIDATOR:
        print(f"  UNEXPECTED validator-site bytes. Expected {ORIG_VALIDATOR.hex()}")
        CloseHandle(h); sys.exit(5)

    print(f"  writing validator ({len(VALIDATOR_BYTES)} bytes)")
    write_bytes(h, VALIDATOR_VA, VALIDATOR_BYTES)
    print(f"    verify: {read_bytes(h, VALIDATOR_VA, 29).hex()}")
    print(f"  writing dispatch redirect")
    write_bytes(h, DISPATCH_VA, NEW_DISPATCH, prot=PAGE_READWRITE)
    print(f"    verify: {read_bytes(h, DISPATCH_VA, 4).hex()}")
    print(f"  OK")
    CloseHandle(h)


if __name__ == "__main__":
    main()