"""position_array_inspect.py <pid> -- for arrays of Positions (96.6%
of all hits), find the bytes immediately BEFORE the array start and
immediately AFTER the array end. This is where the array allocator's
header would live, or where the containing DArray/SmartArray descriptor
points.
"""
import ctypes, ctypes.wintypes as wt, struct, sys
from collections import Counter

POSITION_VT = 0x00797910
SIZEOF_POSITION = 72
ACMIN, ACMAX = 0x00400000, 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 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

    # For each "array head" (Position with no Position 72 bytes before),
    # collect: (array_len, bytes 16 BEFORE first position).
    # For each "array tail" (Position with no Position 72 bytes after),
    # collect: bytes 16 AFTER end of position.
    array_lens = Counter()
    head_minus_4 = Counter()
    head_minus_8 = Counter()
    head_minus_12 = Counter()
    head_minus_16 = Counter()
    array_size_bytes = Counter()
    total_arrays = 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
                pos_offs = []
                for off in range(0, end, 4):
                    if struct.unpack_from("<I", data, off)[0] == POSITION_VT:
                        pos_offs.append(off)
                pos_set = set(pos_offs)
                # find array heads
                for off in pos_offs:
                    if (off - SIZEOF_POSITION) in pos_set:
                        continue  # not a head
                    # Walk forward to find the tail
                    tail = off
                    length = 1
                    while (tail + SIZEOF_POSITION) in pos_set:
                        tail += SIZEOF_POSITION
                        length += 1
                    if length < 2:
                        continue  # not really an array
                    total_arrays += 1
                    array_lens[length] += 1
                    array_size_bytes[length * SIZEOF_POSITION] += 1
                    # Bytes BEFORE head
                    if off >= 16:
                        m4 = struct.unpack_from("<I", data, off - 4)[0]
                        m8 = struct.unpack_from("<I", data, off - 8)[0]
                        m12 = struct.unpack_from("<I", data, off - 12)[0]
                        m16 = struct.unpack_from("<I", data, off - 16)[0]
                        head_minus_4[m4] += 1
                        head_minus_8[m8] += 1
                        head_minus_12[m12] += 1
                        head_minus_16[m16] += 1
        addr = (mbi.BaseAddress or 0) + mbi.RegionSize
        if addr >= 0x80000000:
            break

    print(f"PID {pid}: {total_arrays} Position arrays found")
    print()
    print("Top array lengths (positions per array):")
    for length, n in array_lens.most_common(15):
        print(f"  len={length:>5}  count={n:>6}  bytes/array={length*SIZEOF_POSITION:>7}")
    print()
    print("Top values at offset -4 of array head (likely refcount/size/flags):")
    for v, n in head_minus_4.most_common(8):
        print(f"  {v:#010x}  count={n}")
    print()
    print("Top values at offset -8 of array head:")
    for v, n in head_minus_8.most_common(8):
        print(f"  {v:#010x}  count={n}")
    print()
    print("Top values at offset -12 of array head:")
    for v, n in head_minus_12.most_common(8):
        print(f"  {v:#010x}  count={n}")
    print()
    print("Top values at offset -16 of array head:")
    for v, n in head_minus_16.most_common(8):
        print(f"  {v:#010x}  count={n}")

    k.CloseHandle(h)
    return 0


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