"""scan_gmui_subclasses.py <dump.dmp>

Scan for all aligned dwords equal to 0x007ccb60 (NoticeHandler vtable at
offset 0x5F8 of every gm*UI instance). For each hit at address H, read
the OUTER instance vtable at (H - 0x5F8) — that's the instance's primary
vtable at offset 0. Histogram by outer vtable.

Also dump the SECONDARY vtables at well-known offsets to help disambiguate
subclasses (since gm*UI uses multiple inheritance).
"""
import struct, sys
from collections import Counter, defaultdict
from minidump.minidumpfile import MinidumpFile


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


NOTICE_VT = 0x007ccb60
NOTICE_OFF = 0x5F8


def main():
    md = MinidumpFile.parse(sys.argv[1])
    reader = md.get_reader().get_buffered_reader()

    # Module map
    mods = []
    for m in md.modules.modules:
        mods.append((m.baseaddress, m.size, m.name))
    def mod_of(addr):
        for b, s, n in mods:
            if b <= addr < b + s:
                return n.split("\\")[-1]
        return None

    # Image-region ranges (for vtable validation)
    image_ranges = []
    for r in md.memory_info.infos:
        st, ty = _ei(r.State), _ei(r.Type)
        if st == 0x1000 and ty == 0x1000000:
            image_ranges.append((r.BaseAddress, r.BaseAddress + r.RegionSize))
    image_ranges.sort()
    def is_image(addr):
        for lo, hi in image_ranges:
            if lo <= addr < hi:
                return True
            if addr < lo:
                return False
        return False

    # Scan all committed RW non-image regions for the NoticeHandler vtable value
    scan_regions = []
    for r in md.memory_info.infos:
        st, ty, pr = _ei(r.State), _ei(r.Type), _ei(r.Protect) & 0xff
        if st != 0x1000: continue
        if ty == 0x1000000: continue
        if pr not in (0x04, 0x40): continue
        scan_regions.append((r.BaseAddress, r.RegionSize))

    total_bytes = sum(s for _, s in scan_regions)
    print(f"scanning {len(scan_regions)} writable non-image regions ({total_bytes/(1024*1024):.1f} MB)")

    # All hit addresses (where the NoticeHandler vtable value appears)
    hit_addrs = []
    region_buffers = {}  # base -> (size, bytes)
    for base, size in scan_regions:
        try:
            reader.move(base)
            buf = reader.read(size)
        except Exception:
            continue
        if not buf: continue
        region_buffers[base] = (size, buf)
        end = (len(buf) // 4) * 4
        for off in range(0, end, 4):
            v = struct.unpack_from("<I", buf, off)[0]
            if v == NOTICE_VT:
                hit_addrs.append(base + off)

    print(f"raw NoticeHandler-vtable hits: {len(hit_addrs)}")

    # For each hit, derive instance base = hit - 0x5F8 and try to read the
    # outer vtable at instance_base + 0. Histogram outer vtable values.
    def read_dword(va):
        for base, (size, buf) in region_buffers.items():
            if base <= va < base + size:
                off = va - base
                if off + 4 <= len(buf):
                    return struct.unpack_from("<I", buf, off)[0]
                return None
        return None

    outer_vt_hist = Counter()
    instances = []  # (instance_base, outer_vt)
    not_in_rw = 0
    not_image_vt = 0

    for h in hit_addrs:
        inst = h - NOTICE_OFF
        outer = read_dword(inst)
        if outer is None:
            not_in_rw += 1
            continue
        if not is_image(outer):
            not_image_vt += 1
            continue
        outer_vt_hist[outer] += 1
        instances.append((inst, outer))

    print(f"valid gm*UI instances (outer vtable is in image): {len(instances)}")
    print(f"  hits where instance base not in RW region: {not_in_rw}")
    print(f"  hits where outer vtable is non-image (false positive?): {not_image_vt}")
    print()

    print("=== Outer vtable histogram (top 15) ===")
    for vt, cnt in outer_vt_hist.most_common(15):
        owner = mod_of(vt) or "?"
        print(f"  0x{vt:08x}  count={cnt:<5}  ({owner})")

    # For each top outer vtable, sample a few instances and dump all
    # image-pointer values within the instance (to find secondary vtables).
    print()
    print("=== Secondary vtable scan (top outer vtables) ===")
    # Group instances by outer vtable
    by_outer = defaultdict(list)
    for inst, outer in instances:
        by_outer[outer].append(inst)

    SCAN_SIZE = 0x600  # scan 1.5KB of instance (covers up to NoticeHandler @ 0x5F8)
    for outer, _ in outer_vt_hist.most_common(8):
        insts = by_outer[outer]
        print(f"\n-- outer vtable 0x{outer:08x}  ({len(insts)} instances) --")

        # For each offset in the instance, count how many instances have an
        # image-pointer at that offset, and what the most common value is.
        offset_value_hist = defaultdict(Counter)
        sample_n = min(50, len(insts))
        for inst in insts[:sample_n]:
            for off in range(0, SCAN_SIZE, 4):
                v = read_dword(inst + off)
                if v is None: continue
                if is_image(v):
                    offset_value_hist[off][v] += 1

        # Show offsets where the value is consistent across most samples
        consistent_offsets = []
        for off, vh in offset_value_hist.items():
            top_v, top_c = vh.most_common(1)[0]
            if top_c >= sample_n * 0.7:  # 70% agreement
                consistent_offsets.append((off, top_v, top_c, len(vh)))
        consistent_offsets.sort()

        print(f"  consistent image-pointers (>=70% of {sample_n} samples):")
        for off, val, cnt, unique in consistent_offsets:
            tag = ""
            if off == 0: tag = " [PRIMARY VTABLE]"
            elif off == NOTICE_OFF: tag = " [NoticeHandler]"
            print(f"    +0x{off:04x}  val=0x{val:08x}  ({cnt}/{sample_n}, {unique} unique){tag}")


if __name__ == "__main__":
    main()