#!/usr/bin/env python3
"""Exhaustive reference finder for the RT-4D app image.

Decodes one instruction at EVERY 2-byte alignment (linear sweep desyncs on
Thumb), so nothing is missed.

  python refs.py addr <0xADDR> [ctx]   who pc-relative-loads this address/constant
  python refs.py gpio                  full GPIO map: (port, register, mask) per site
"""
import sys, struct, re, capstone

BASE = 0x08002800
IMG = open("/home/viktor/claude/rt-4d/stock-fw/rt4d_stock_v3.25_abs_0x08002800.bin", "rb").read()
md = capstone.Cs(capstone.CS_ARCH_ARM, capstone.CS_MODE_THUMB)

PORTS = {0x40020000: "GPIOA", 0x40020400: "GPIOB", 0x40020800: "GPIOC",
         0x40020C00: "GPIOD", 0x40021000: "GPIOE", 0x40021400: "GPIOF"}
# AT32F43x GPIO layout
REG = {0x00: "CFGR", 0x10: "IDT", 0x14: "ODT", 0x18: "SCR(set)",
       0x1C: "CLR16", 0x28: "CLR(clear)", 0x2C: "TOGR(toggle)"}


def word(va):
    o = va - BASE
    return struct.unpack_from("<I", IMG, o)[0] if 0 <= o <= len(IMG) - 4 else None


def decode_at(va):
    o = va - BASE
    for i in md.disasm(IMG[o:o + 4], va):
        return i
    return None


def litval(i):
    if i is None or not i.mnemonic.startswith("ldr"):
        return None
    m = re.search(r"\[pc, #(?:0x)?([0-9a-fA-F]+)\]", i.op_str)
    if not m:
        return None
    return word(((i.address + 4) & ~3) + int(m.group(1), 16))


def all_lit_loads():
    """[(va, ins, value)] for every pc-relative ldr at any 2-byte alignment"""
    out = []
    for off in range(0, len(IMG) - 4, 2):
        va = BASE + off
        i = decode_at(va)
        v = litval(i)
        if v is not None:
            out.append((va, i, v))
    return out


def cmd_addr(target, ctx=8):
    hits = [(va, i) for va, i, v in all_lit_loads() if v == target]
    print(f"pc-relative loads of 0x{target:08X}: {len(hits)}")
    for va, i in hits:
        print(f"\n--- 0x{va:08X}  ({i.mnemonic} {i.op_str}) ---")
        a = va
        for _ in range(ctx):
            j = decode_at(a)
            if j is None:
                break
            v = litval(j)
            e = f"   ; =0x{v:08X}" if v is not None else ""
            print(f"   0x{j.address:08X}: {j.mnemonic:8} {j.op_str}{e}")
            a += j.size


def cmd_gpio():
    loads = all_lit_loads()
    rows = []
    for va, i, v in loads:
        if v not in PORTS:
            continue
        reg = i.op_str.split(",")[0].strip()
        a = va + i.size
        mask = None
        # small forward window: find store to [reg,#off]
        for _ in range(6):
            j = decode_at(a)
            if j is None:
                break
            m = re.match(r"(r\d+), \[" + reg + r", #(0x[0-9a-fA-F]+|\d+)\]$", j.op_str)
            if j.mnemonic.startswith("str") and m:
                src, off = m.group(1), int(m.group(2), 0)
                # back-window for the immediate loaded into src
                b = va
                for _ in range(10):
                    b -= 2
                    p = decode_at(b)
                    if p and p.op_str.startswith(src + ", #"):
                        try:
                            mask = int(p.op_str.split("#")[1], 0)
                        except Exception:
                            pass
                        break
                rows.append((va, PORTS[v], off, mask))
                break
            a += j.size
    print(f"{len(rows)} GPIO store sites\n")
    print(f"{'site':<13}{'port':<7}{'register':<14}{'mask':<10}bit")
    for va, p, off, m in rows:
        bit = f"bit{m.bit_length()-1}" if m and (m & (m - 1)) == 0 else ("?" if m is None else "multi")
        print(f"0x{va:08X}  {p:<7}{REG.get(off, hex(off)):<14}{(hex(m) if m is not None else '?'):<10}{bit}")


if __name__ == "__main__":
    if sys.argv[1] == "addr":
        cmd_addr(int(sys.argv[2], 16), int(sys.argv[3]) if len(sys.argv) > 3 else 8)
    elif sys.argv[1] == "gpio":
        cmd_gpio()