"""identify_mystery_vtables.py <pid> <vtable1> [vtable2 ...]

For each vtable address, dump the first 16 function pointers and try to
match them against references/symbols.json for class identification.
"""
import argparse, ctypes, ctypes.wintypes as wt, json, struct, sys, os

ap = argparse.ArgumentParser()
ap.add_argument("pid", type=int)
ap.add_argument("vtables", nargs="+")
args = ap.parse_args()

vtables = [int(v, 0) for v in args.vtables]

PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400

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

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

def read(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])

# Load symbols
syms = {}
sym_path = os.path.join(os.path.dirname(__file__), "..", "references", "symbols.json")
if os.path.exists(sym_path):
    with open(sym_path) as f:
        for s in json.load(f):
            try:
                a = int(s["address"], 16)
                syms[a] = s.get("name", "?")
            except Exception:
                pass
print(f"Loaded {len(syms)} symbols")

# Note: these are EoR addresses; we don't have an EoR symbols file.
# But we may still see overlapping addresses (low frequency) — or near-match heuristic.
# Better: just dump and let the human interpret. Also dump strings near each
# class' first method (often vftable strings live around the .rdata block).

for vt in vtables:
    print(f"\n=== vtable 0x{vt:08x} ===")
    data = read(vt, 64)
    if not data:
        print("  unreadable"); continue
    for i in range(16):
        p = struct.unpack_from("<I", data, i*4)[0]
        name = syms.get(p, "")
        # Also probe for proximity matches (within 0x40)
        near = ""
        if not name:
            for d in range(0, 0x40, 1):
                if (p - d) in syms:
                    near = f"  ~ {syms[p-d]}+0x{d:x}"; break
        print(f"  [{i:2d}] 0x{p:08x}  {name}{near}")

    # Also look around the vtable - 4 (often points to RTTI structure)
    rtti_slot = read(vt - 4, 4)
    if rtti_slot:
        rtti_ptr = struct.unpack("<I", rtti_slot)[0]
        print(f"  RTTI slot (vt-4): 0x{rtti_ptr:08x}")
        # Read RTTI COL: typeDescriptor at +0xC
        col = read(rtti_ptr, 0x20)
        if col:
            type_desc = struct.unpack_from("<I", col, 0xC)[0]
            print(f"  RTTI typeDescriptor: 0x{type_desc:08x}")
            # Type descriptor: vtable(4) + spare(4) + name(zero-terminated str)
            name_addr = type_desc + 8
            name_bytes = read(name_addr, 128)
            if name_bytes:
                end = name_bytes.find(b'\x00')
                if end > 0:
                    try:
                        decoded = name_bytes[:end].decode("latin-1")
                        print(f"  RTTI name: {decoded!r}")
                    except Exception:
                        print(f"  RTTI raw: {name_bytes[:end]!r}")