"""
scan_vtables.py <dump.dmp> [--size-min KB] [--size-max KB]

For every private RW region in [size-min, size-max] (default 256-512 KB),
read the first 0x100 bytes and scan for DWORDs that look like vtable
pointers (addresses pointing into acclient.exe's code/data range).

Outputs a histogram of (vtable RVA -> count of regions where it appeared
at any of the first 64 DWORD positions). The top entries are candidate
leaking classes.

Why this works without symbols:
  * 32-bit C++ objects start with a vtable pointer at offset 0.
  * The NT heap large-block path wraps a small header around the user
    data, so the vtable shows up somewhere in the first 0x80 bytes.
  * If thousands of leaked regions all carry the same vtable, it is
    the same class — the RVA can be looked up in the EoR binary
    (relative to acclient.exe base 0x00400000 typically) for an AOB
    signature later.
"""
import argparse
import os
import sys
from collections import Counter

from minidump.minidumpfile import MinidumpFile


def _enum_int(v):
    if v is None: return 0
    if hasattr(v, 'value'): return int(v.value)
    return int(v)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dump")
    ap.add_argument("--size-min-kb", type=int, default=256)
    ap.add_argument("--size-max-kb", type=int, default=512)
    ap.add_argument("--scan-bytes",  type=int, default=0x100,
                    help="bytes from region start to scan for vtables")
    ap.add_argument("--top",         type=int, default=20)
    args = ap.parse_args()

    md = MinidumpFile.parse(args.dump)

    # Find acclient.exe base + size
    acl = None
    for m in md.modules.modules:
        if os.path.basename(m.name).lower() == "acclient.exe":
            acl = m
            break
    if acl is None:
        print("acclient.exe not in module list", file=sys.stderr); sys.exit(1)
    acl_lo = acl.baseaddress
    acl_hi = acl.baseaddress + acl.size
    print(f"acclient.exe: 0x{acl_lo:08x} - 0x{acl_hi:08x}  ({acl.size/1024:.1f} KB)")

    # Build a reader. minidump has a memory reader.
    reader = md.get_reader().get_buffered_reader()

    # Iterate candidate regions
    lo_b = args.size_min_kb * 1024
    hi_b = args.size_max_kb * 1024
    cands = []
    for r in md.memory_info.infos:
        st = _enum_int(r.State)
        ty = _enum_int(r.Type)
        pr = _enum_int(r.Protect) & 0xFF
        if st == 0x1000 and ty == 0x20000 and pr in (0x04, 0x40) \
                and lo_b <= r.RegionSize < hi_b:
            cands.append((r.BaseAddress, r.RegionSize))
    print(f"candidate regions ({args.size_min_kb}-{args.size_max_kb}KB private RW): {len(cands)}")

    # For each region: read first scan_bytes, scan as DWORDs
    hits = Counter()           # vtable_rva -> region_count
    region_hits = Counter()    # how many regions had ANY hit
    per_region_seen = []
    total_dword_hits = 0

    for base, size in cands:
        try:
            reader.move(base)
            buf = reader.read(args.scan_bytes)
        except Exception:
            continue
        if not buf:
            continue
        seen_here = set()
        for off in range(0, min(len(buf), args.scan_bytes), 4):
            dw = int.from_bytes(buf[off:off+4], "little", signed=False)
            if acl_lo <= dw < acl_hi:
                rva = dw - acl_lo
                # An object's vtable typically sits in .rdata so RVAs are
                # mid-image. We don't filter further here; let the histogram
                # speak.
                seen_here.add(rva)
                total_dword_hits += 1
        for rva in seen_here:
            hits[rva] += 1
        if seen_here:
            region_hits[len(seen_here)] += 1
            per_region_seen.append(len(seen_here))

    print(f"regions with at least one acclient-pointer in first {args.scan_bytes:#x} bytes: {sum(region_hits.values())}")
    print(f"distribution of pointer-counts per region:")
    for n in sorted(region_hits):
        print(f"  {n} pointer(s) in region: {region_hits[n]} regions")
    print()
    print(f"top {args.top} vtable-candidate RVAs (regions sharing this acclient pointer in first {args.scan_bytes:#x} bytes):")
    print(f"  {'rva':>10}  {'abs':>10}  {'count':>6}")
    for rva, count in hits.most_common(args.top):
        print(f"  0x{rva:08x}  0x{acl_lo+rva:08x}  {count:>6}")


if __name__ == "__main__":
    main()