"""find_mesh_refs_inc_static.py <pid>

Like trace_mesh_holder but ALSO scans MEM_IMAGE regions (e.g., .data
section of acclient.exe) for references to D3DXMesh instances.
"""
import ctypes, ctypes.wintypes as wt, struct, sys

VTABLE = 0x007ed3b0
PROCESS_VM_READ=0x10; PROCESS_QUERY_INFORMATION=0x400
MEM_COMMIT=0x1000; MEM_PRIVATE=0x20000; MEM_IMAGE=0x1000000

class MBI(ctypes.Structure):
    _fields_ = [('BaseAddress',ctypes.c_void_p),('AllocationBase',ctypes.c_void_p),
                ('AllocationProtect',wt.DWORD),('PartitionId',wt.WORD),('RegionSize',ctypes.c_size_t),
                ('State',wt.DWORD),('Protect',wt.DWORD),('Type',wt.DWORD)]
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,ctypes.c_void_p,ctypes.POINTER(MBI),ctypes.c_size_t]
k.VirtualQueryEx.restype=ctypes.c_size_t

pid = int(sys.argv[1])
h = k.OpenProcess(PROCESS_VM_READ|PROCESS_QUERY_INFORMATION, False, pid)

all_regions = []  # tuples (base, data, kind)
mbi=MBI(); addr=0
while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    pr = mbi.Protect & 0xff
    if mbi.State==MEM_COMMIT and pr != 0x01:
        kind = 'priv' if mbi.Type==MEM_PRIVATE else ('image' if mbi.Type==MEM_IMAGE else 'other')
        # Read writable regions of EITHER private or image
        if pr in (0x04, 0x40, 0x02):  # RW, ERW, or RO
            buf=(ctypes.c_ubyte*mbi.RegionSize)(); sz=ctypes.c_size_t(0)
            if k.ReadProcessMemory(h, mbi.BaseAddress, buf, mbi.RegionSize, ctypes.byref(sz)):
                all_regions.append((mbi.BaseAddress, bytes(buf[:sz.value]), kind))
    addr=(mbi.BaseAddress or 0)+mbi.RegionSize
    if addr>=0x80000000: break

# Find mesh instance addresses (only in heap RW; not in image)
mesh_addrs = set()
for base, data, kind in all_regions:
    if kind != 'priv': continue
    end = (len(data)//4)*4
    for off in range(0, end-4, 4):
        if struct.unpack_from('<I', data, off)[0] == VTABLE:
            mesh_addrs.add(base + off)
print(f'mesh instances found: {len(mesh_addrs)}')

# Count references to each mesh, distinguishing heap vs image
heap_refs = {a: 0 for a in mesh_addrs}
image_refs = {a: 0 for a in mesh_addrs}
for base, data, kind in all_regions:
    end = (len(data)//4)*4
    for off in range(0, end-4, 4):
        v = struct.unpack_from('<I', data, off)[0]
        if v in mesh_addrs:
            ra = base + off
            if ra in mesh_addrs: continue
            if kind == 'priv':
                heap_refs[v] += 1
            elif kind == 'image':
                image_refs[v] += 1

# Classify
true_orphans = 0       # 0 heap + 0 image refs
heap_held = 0          # has at least 1 heap ref
image_only = 0         # 0 heap but has image ref(s)
for a in mesh_addrs:
    if heap_refs[a] == 0 and image_refs[a] == 0: true_orphans += 1
    elif heap_refs[a] > 0: heap_held += 1
    else: image_only += 1

print(f'true orphans (0 heap + 0 image refs): {true_orphans}')
print(f'heap-held: {heap_held}')
print(f'image-only refs (static globals): {image_only}')

# For image-only refs, show some examples (the static global location)
print()
print('Sample image-only refs:')
shown = 0
for a in mesh_addrs:
    if heap_refs[a] == 0 and image_refs[a] > 0 and shown < 10:
        # find the image ref locations
        for base, data, kind in all_regions:
            if kind != 'image': continue
            end = (len(data)//4)*4
            for off in range(0, end-4, 4):
                if struct.unpack_from('<I', data, off)[0] == a:
                    print(f'  mesh @ 0x{a:08x} <- image ref @ 0x{base+off:08x}')
                    shown += 1
                    if shown >= 10: break
            if shown >= 10: break