#!/usr/bin/env python3
"""Scan MediaTek Connac2 firmware regions for {id, fn*} registration arrays.

WM registers command handlers as sparse {u8 cmd_id, code*} arrays (see
findings F9: dispatcher walks gp+0x1b6a0, stride 8). This tool finds every
such array across a blob's regions — the recoverable handler map.

Usage: scan_registrations.py <blob-stem-dir> [min_run] > out.txt
(blob-stem-dir = extracted/mt7981_wm; scans all rN region files)
"""
import struct
import sys
from pathlib import Path


def code_p(w, ranges):
    return any(lo <= w < hi for lo, hi in ranges)


def scan(f: Path, ranges, min_run):
    d = f.read_bytes()
    base = int(f.name.split('_a')[1][:8], 16)
    n = len(d) // 4
    ws = struct.unpack_from('<%dI' % n, d)
    out = []
    i = 0
    while i < n - 1:
        j, run = i, 0
        while j < n - 1 and ws[j] < 0x100 and code_p(ws[j + 1], ranges):
            run += 1
            j += 2
        if run >= min_run:
            out.append((base + i * 4, run, ws[i:j]))
        i = j + 2 if run else i + 1
    return out


def main():
    stem = Path(sys.argv[1])
    min_run = int(sys.argv[2]) if len(sys.argv) > 2 else 8
    ranges = []
    for f in sorted(stem.glob('r*.bin')):
        base = int(f.name.split('_a')[1][:8], 16)
        ranges.append((base, base + f.stat().st_size))
    for f in sorted(stem.glob('r*.bin')):
        for addr, run, ws in scan(f, ranges, min_run):
            print('# %d-entry {id,fn} array @ 0x%08x (%s)' % (run, addr, f.name))
            for k in range(run):
                print('0x%08x id=0x%02x fn=0x%08x' % (addr + k * 8, ws[2 * k], ws[2 * k + 1]))


if __name__ == '__main__':
    main()