"""position_heap_solo_scan.py <pid>

Variant of position_host_scan that filters to ONLY Positions which
look like standalone heap allocations (vs embedded into a bigger
struct).

Heuristic for "standalone heap allocation":
  - Bytes right before the Position vtable have a typical heap-block
    header pattern (small DWORD = allocation size, then a possibly-NULL
    block-id, then padding/flags).
  - OR there is a ref-counted wrapper vtable at offset -8 (the
    PositionPropertyValue pattern).
  - The +96..+128 region after the Position is either NULL/heap-free
    (suggesting nothing follows the Position in the same block).

Goal: classify Positions into:
  (A) embedded in a larger object — host vtable found at offset -4
      through -16 (the Position is just a field of e.g. a PhysicsObj
      or NetBlob).
  (B) inside a ref-counted property wrapper — vtable at -8
      (PositionPropertyValue layout).
  (C) heap-solo at offset 0 — allocated as `new Position(...)`
      directly, no enclosing object.
  (D) part of an array — another Position vtable 96 bytes away.

The leaking class is whichever bucket dominates the delta between
heavy-looter and idle character.
"""
import ctypes
import ctypes.wintypes as wt
import struct
import sys
from collections import Counter, defaultdict

POSITION_VT = 0x00797910
ACCLIENT_MIN = 0x00400000
ACCLIENT_MAX = 0x00900000

PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400


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 classify(data, off):
    """Return (bucket, evidence_value). bucket in {'array','ppv',
    'embedded','solo','unknown'}."""
    # array: Position vtable 96 bytes earlier or later
    if off >= 96:
        prev = struct.unpack_from("<I", data, off - 96)[0]
        if prev == POSITION_VT:
            return ('array_prev', prev)
    if off + 96 + 4 <= len(data):
        nxt = struct.unpack_from("<I", data, off + 96)[0]
        if nxt == POSITION_VT:
            return ('array_next', nxt)
    # ppv: vtable at off-8, refcount-shaped dword at off-4
    if off >= 8:
        m8 = struct.unpack_from("<I", data, off - 8)[0]
        m4 = struct.unpack_from("<I", data, off - 4)[0]
        if ACCLIENT_MIN <= m8 <= ACCLIENT_MAX and 0 < m4 < 0x100000:
            return ('ppv', m8)
    # embedded: any acclient code-range vtable at offsets -32..-4
    if off >= 32:
        for d in range(-32, 0, 4):
            v = struct.unpack_from("<I", data, off + d)[0]
            if ACCLIENT_MIN <= v <= ACCLIENT_MAX:
                return ('embedded', v)
    # solo: nothing identifying
    return ('solo', 0)


def scan(pid):
    h = k.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
    if not h:
        print(f"OpenProcess err={ctypes.get_last_error()}"); return 1

    bucket_counts = Counter()
    embedded_vt_counts = Counter()
    ppv_vt_counts = Counter()
    array_off_counts = Counter()

    total = 0
    mbi = MBI()
    addr = 0
    while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
        pr = mbi.Protect & 0xff
        if (mbi.State == 0x1000 and mbi.Type == 0x20000
                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])
                end = (len(data) // 4) * 4
                for off in range(0, end, 4):
                    if struct.unpack_from("<I", data, off)[0] == POSITION_VT:
                        total += 1
                        b, ev = classify(data, off)
                        bucket_counts[b] += 1
                        if b == 'embedded':
                            embedded_vt_counts[ev] += 1
                        elif b == 'ppv':
                            ppv_vt_counts[ev] += 1
                        elif b.startswith('array'):
                            array_off_counts[b] += 1
        addr = (mbi.BaseAddress or 0) + mbi.RegionSize
        if addr >= 0x80000000:
            break

    print(f"PID {pid}: {total} Position instances scanned")
    print()
    print("Bucket distribution:")
    for b, n in bucket_counts.most_common():
        pct = 100.0 * n / max(1, total)
        print(f"  {b:>12}: {n:>7}  ({pct:5.1f}%)")
    print()
    print("Top 'embedded' host vtables (Position is a member of class):")
    for v, n in embedded_vt_counts.most_common(15):
        print(f"  0x{v:08x}  count={n:>6}")
    print()
    print("Top 'ppv' wrapper vtables (Position inside ref-counted holder):")
    for v, n in ppv_vt_counts.most_common(10):
        print(f"  0x{v:08x}  count={n:>6}")
    print()
    print("Array adjacency:")
    for b, n in array_off_counts.most_common():
        print(f"  {b}: {n}")

    k.CloseHandle(h)
    return 0


if __name__ == "__main__":
    sys.exit(scan(int(sys.argv[1])))