"""count_palettes_live.py <pid>
Count Palette instances (vtable 0x007caa08) in a live process and
break down by refcount.
"""
import ctypes, ctypes.wintypes as wt, struct, sys
from collections import Counter

VTABLE = 0x007caa08
PROCESS_VM_READ = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400
MEM_COMMIT = 0x1000
MEM_PRIVATE = 0x20000

class MEMORY_BASIC_INFORMATION(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),
    ]

k32 = ctypes.windll.kernel32
OpenProcess = k32.OpenProcess
OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]; OpenProcess.restype = wt.HANDLE
CloseHandle = k32.CloseHandle
CloseHandle.argtypes = [wt.HANDLE]; CloseHandle.restype = wt.BOOL
ReadProcessMemory = k32.ReadProcessMemory
ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t,
                              ctypes.POINTER(ctypes.c_size_t)]
ReadProcessMemory.restype = wt.BOOL
VirtualQueryEx = k32.VirtualQueryEx
VirtualQueryEx.argtypes = [wt.HANDLE, ctypes.c_void_p,
                           ctypes.POINTER(MEMORY_BASIC_INFORMATION), ctypes.c_size_t]
VirtualQueryEx.restype = ctypes.c_size_t


pid = int(sys.argv[1])
h = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
if not h:
    print(f"OpenProcess({pid}) err={ctypes.get_last_error()}"); sys.exit(2)

mbi = MEMORY_BASIC_INFORMATION()
addr = 0
n_total = 0
rc_hist = Counter()
maintainer_hist = Counter()
m_numlinks_hist = Counter()

while VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    base = mbi.BaseAddress or 0
    size = mbi.RegionSize
    st = mbi.State
    ty = mbi.Type
    pr = mbi.Protect & 0xFF
    if st == MEM_COMMIT and ty == MEM_PRIVATE and pr in (0x04, 0x40):
        buf = (ctypes.c_ubyte * size)()
        sz = ctypes.c_size_t(0)
        if ReadProcessMemory(h, base, buf, size, ctypes.byref(sz)) and sz.value > 0x48:
            data = bytes(buf[:sz.value])
            end = (len(data) // 4) * 4
            for off in range(0, end - 0x48, 4):
                if struct.unpack_from("<I", data, off)[0] == VTABLE:
                    rc = struct.unpack_from("<I", data, off + 0x24)[0]
                    main = struct.unpack_from("<I", data, off + 0x20)[0]
                    numl = struct.unpack_from("<I", data, off + 0x04)[0]
                    rc_hist[rc if rc < 100 else 100] += 1
                    maintainer_hist[1 if main else 0] += 1
                    m_numlinks_hist[numl & 0xff] += 1
                    n_total += 1
    addr = base + size
    if addr >= 0x80000000: break

print(f"PID {pid}: {n_total} Palette instances")
print("refcount distribution (top 8):")
for rc, n in rc_hist.most_common(8):
    print(f"  rc={rc:<4}  {n}")
print(f"m_pMaintainer NULL: {maintainer_hist[0]}, non-NULL: {maintainer_hist[1]}")
print(f"m_numLinks distribution (top 6):")
for ml, n in m_numlinks_hist.most_common(6):
    print(f"  ml={ml:<4}  {n}")

CloseHandle(h)