"""count_weenie_live.py <pid> [<pid> ...]
Count ACCWeenieObject + CWeenieObject vtable instances in committed-private RW
heap regions of one or more live AC clients."""
import ctypes, ctypes.wintypes as wt, struct, sys

VTABLES = {
    "ACCWeenieObject_pri":   0x007e4f70,
    "ACCWeenieObject_sec":   0x007e4f58,
    "CWeenieObject_pri":     0x007e4ed8,
    "CWeenieObject_sec":     0x007e4ec4,
}

PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400
MEM_COMMIT = 0x1000
MEM_PRIVATE = 0x20000

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.CloseHandle.argtypes = [wt.HANDLE]; k.CloseHandle.restype = wt.BOOL
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

def scan(pid):
    h = k.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
    if not h:
        print(f"OpenProcess({pid}) err={ctypes.get_last_error()}"); return None
    counts = {n: 0 for n in VTABLES}
    vt_to_name = {vt: name for name, vt in VTABLES.items()}
    mbi = MBI()
    addr = 0
    total_bytes = 0
    while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
        pr = mbi.Protect & 0xff
        if (mbi.State == MEM_COMMIT and mbi.Type == MEM_PRIVATE and pr in (0x04, 0x40)):
            buf = (ctypes.c_ubyte * mbi.RegionSize)()
            sz = ctypes.c_size_t(0)
            if k.ReadProcessMemory(h, mbi.BaseAddress, buf, mbi.RegionSize, ctypes.byref(sz)):
                data = bytes(buf[:sz.value])
                total_bytes += sz.value
                end = (len(data) // 4) * 4
                for off in range(0, end, 4):
                    v = struct.unpack_from("<I", data, off)[0]
                    if v in vt_to_name:
                        counts[vt_to_name[v]] += 1
        addr = (mbi.BaseAddress or 0) + mbi.RegionSize
        if addr >= 0x80000000: break
    k.CloseHandle(h)
    return counts, total_bytes

if __name__ == "__main__":
    print(f"{'pid':>6}  {'ACC_pri':>8} {'ACC_sec':>8} {'CW_pri':>8} {'CW_sec':>8}   priv_MB")
    for arg in sys.argv[1:]:
        pid = int(arg)
        r = scan(pid)
        if r is None: continue
        c, b = r
        print(f"{pid:6d}  {c['ACCWeenieObject_pri']:8d} {c['ACCWeenieObject_sec']:8d} "
              f"{c['CWeenieObject_pri']:8d} {c['CWeenieObject_sec']:8d}   {b/1024/1024:7.1f}")