"""find_d3d_device.py <pid>
Find AC's live IDirect3DDevice9 instance.
1. Enumerate loaded modules; locate d3d9.dll's address range
2. Walk private RW memory; find any pointer where:
   - The pointer itself is in private RW memory (object on heap)
   - The first DWORD (vtable) lies in d3d9.dll's read-only data range
3. Report the first few matches. The one used most often is THE device."""
import ctypes, ctypes.wintypes as wt, sys, struct, collections

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
k.VirtualQueryEx.argtypes = [wt.HANDLE, wt.LPCVOID, ctypes.c_void_p, ctypes.c_size_t]
k.VirtualQueryEx.restype = ctypes.c_size_t

class MBI(ctypes.Structure):
    _fields_ = [("BaseAddress", ctypes.c_void_p), ("AllocationBase", ctypes.c_void_p),
                ("AllocationProtect", wt.DWORD), ("RegionSize", ctypes.c_size_t),
                ("State", wt.DWORD), ("Protect", wt.DWORD), ("Type", wt.DWORD)]

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

def find_d3d9_range(h):
    """Find d3d9.dll's load range by looking at image-mapped regions and
    reading their PE headers to identify the module name."""
    candidates = []
    mbi = MBI()
    addr = 0
    seen_bases = set()
    while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
        base = mbi.BaseAddress or 0
        sz = mbi.RegionSize
        if mbi.State == 0x1000 and mbi.Type == 0x1000000:  # MEM_IMAGE
            ab = mbi.AllocationBase or 0
            if ab not in seen_bases:
                seen_bases.add(ab)
                # Read PE header to identify the dll
                pe_off_b = rd(h, ab + 0x3C, 4)
                if pe_off_b and len(pe_off_b) == 4:
                    pe_off = struct.unpack('<I', pe_off_b)[0]
                    if pe_off < 0x1000:
                        # Read SizeOfImage from optional header
                        sz_b = rd(h, ab + pe_off + 4 + 20 + 56, 4)
                        if sz_b:
                            img_size = struct.unpack('<I', sz_b)[0]
                            candidates.append((ab, img_size))
        next_addr = base + sz
        if next_addr <= addr: break
        addr = next_addr
        if addr >= 0x80000000: break
    # For each candidate image, look for the module name in its export table.
    # Cheaper: scan first 1KB for "d3d9" string.
    for ab, sz in candidates:
        chunk = rd(h, ab, 4096)
        if chunk and b'd3d9' in chunk.lower():
            return ab, sz
    return None, None

pid = int(sys.argv[1])
h = k.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
if not h: print("OpenProcess fail"); sys.exit(2)

d3d9_base, d3d9_size = find_d3d9_range(h)
print(f"d3d9.dll: base=0x{d3d9_base:08x} size={d3d9_size}" if d3d9_base else "d3d9.dll not found")
if not d3d9_base:
    k.CloseHandle(h); sys.exit(0)

# Walk private RW; for each region, find DWORDs in d3d9's range; the addresses
# THEMSELVES (where we found the DWORD) are candidates for "is a vtable pointer".
# But we want the OBJECT addresses — i.e., the location holding the pointer is +0
# of the object. So we report `holding_va` (= object address) where the pointer = vtable.
vtable_hits = collections.Counter()  # vtable_va -> count of pointers to it found
sample_holders = collections.defaultdict(list)

mbi = MBI()
addr = 0
while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    base = mbi.BaseAddress or 0
    sz = mbi.RegionSize
    if (mbi.State == 0x1000 and mbi.Type == 0x20000 and (mbi.Protect & 0xFF) in (0x04, 0x40)
        and sz < 32*1024*1024):
        data = rd(h, base, sz)
        if data:
            n = len(data) // 4
            for i in range(n):
                v = struct.unpack_from('<I', data, i*4)[0]
                if d3d9_base <= v < d3d9_base + d3d9_size:
                    # Looks like a pointer into d3d9 — IF v is itself a vtable
                    # (rather than mid-code), the holder is a COM object.
                    # We're interested in cases where this dword is at offset 0
                    # of an aligned object. For now, just count vtable popularity.
                    vtable_hits[v] += 1
                    holder = base + i*4
                    if len(sample_holders[v]) < 3:
                        sample_holders[v].append(holder)
    next_addr = base + sz
    if next_addr <= addr: break
    addr = next_addr
    if addr >= 0x80000000: break

k.CloseHandle(h)

print(f"\nTop 15 d3d9.dll pointer destinations (likely vtables):")
print(f"  {'vtable':>10} {'count':>6}  sample holders")
for vt, count in vtable_hits.most_common(15):
    holders_s = ' '.join(f'0x{h_:08x}' for h_ in sample_holders[vt])
    print(f"  0x{vt:08x} {count:>6}  {holders_s}")