"""find_d3d9_via_iat.py <pid>
Find d3d9.dll's load address by walking AC's import table.
Strategy: locate the string "Direct3DCreate9" in AC's image, then find
the IAT entry that resolves it. The IAT slot holds the runtime address
of Direct3DCreate9 (= address INSIDE d3d9.dll). VirtualQuery on that
address gives us d3d9.dll's base + size."""
import ctypes, ctypes.wintypes as wt, sys, struct

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])

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)

# AC's image starts at 0x00400000. Read its full size from PE header.
hdr = rd(h, 0x00400000 + 0x3C, 4)
if not hdr: print("can't read AC PE header"); sys.exit(2)
pe_off = struct.unpack('<I', hdr)[0]
img_sz = struct.unpack('<I', rd(h, 0x00400000 + pe_off + 4 + 20 + 56, 4))[0]
print(f"AC image: 0x00400000 .. 0x{0x00400000+img_sz:08x} ({img_sz/1024/1024:.1f} MB)")

# Read entire AC image
ac = rd(h, 0x00400000, img_sz)
if not ac: print("can't read AC image"); sys.exit(2)

# Find import directory: data dir entry 1
data_dir_off = pe_off + 4 + 20 + 96 + 8  # 8 = entry 1 * 8 bytes per entry
import_rva, import_size = struct.unpack_from('<II', ac, data_dir_off)
print(f"Import dir RVA=0x{import_rva:x} size={import_size}")

# Walk import descriptors looking for d3d9.dll
off = import_rva
n_desc = 0
while off + 20 <= len(ac):
    olt, ts, fc, name_rva, fthunk = struct.unpack_from('<IIIII', ac, off)
    if olt == 0 and name_rva == 0 and fthunk == 0:
        print(f"end of import desc at #{n_desc} (off=0x{off:x})")
        break
    if name_rva == 0:
        off += 20; n_desc += 1
        continue
    name = ac[name_rva:name_rva+32].split(b'\x00', 1)[0].decode(errors='replace')
    print(f"  desc[{n_desc}] dll={name!r}  IAT-RVA=0x{fthunk:x}")
    n_desc += 1
    if 'd3d9' in name.lower():
        print(f"\nFound import for {name}")
        print(f"  OriginalFirstThunk RVA=0x{olt:x}")
        print(f"  FirstThunk (IAT)    RVA=0x{fthunk:x}")
        # Walk OLT to find Direct3DCreate9 index
        oft = olt or fthunk
        idx = 0
        while oft + 4 <= len(ac):
            entry = struct.unpack_from('<I', ac, oft)[0]
            if entry == 0: break
            if entry & 0x80000000:
                func_name = f"<ord {entry & 0xFFFF}>"
            else:
                # entry is an RVA to a Hint/Name struct: hint(2) + name(asciiz)
                func_name = ac[entry+2:entry+64].split(b'\x00', 1)[0].decode(errors='replace')
            # Read the IAT slot value (runtime address)
            iat_slot_va = 0x00400000 + fthunk + idx*4
            iat_val = struct.unpack('<I', rd(h, iat_slot_va, 4))[0]
            print(f"    [{idx:2d}] {func_name:30s}  IAT@0x{iat_slot_va:08x} -> 0x{iat_val:08x}")
            if func_name == 'Direct3DCreate9':
                # Use this to find d3d9.dll's range
                mbi = MBI()
                if k.VirtualQueryEx(h, iat_val, ctypes.byref(mbi), ctypes.sizeof(mbi)):
                    ab = mbi.AllocationBase or 0
                    # Get image size from PE
                    hdr2 = rd(h, ab + 0x3C, 4)
                    pe2 = struct.unpack('<I', hdr2)[0]
                    sz2 = struct.unpack('<I', rd(h, ab + pe2 + 4 + 20 + 56, 4))[0]
                    print(f"\n*** d3d9.dll loaded at 0x{ab:08x}, size {sz2}")
            oft += 4
            idx += 1
        break
    off += 20
else:
    print("walk fell off")

k.CloseHandle(h)