#!/usr/bin/env python3
"""Annotate a TA in a live idatui session from the decompiler's own switch output.

`fw/ta_symbols.py` recovers entry points statically, but its jump-table walk misses
dispatchers reached by a `b.w` tail call and two-tier switches. The decompiler already
resolves both, so this reads `case N: ... = sub_XXXX(...)` straight out of the
pseudocode and renames each handler `<prefix>_cmdNN`.

Run against a pane already opened on the TA (socket auto-resolved when only one is live).

usage: annotate_ta.py --prefix gcp --dispatch gcp_dispatch [--note "..."]
"""
import argparse
import json
import os
import re
import subprocess
import sys

DRIVE = [sys.executable, '-m', 'idatui.drive']
CWD = os.path.expanduser('~/dev/ida-tui-maybe')


def drive(*args, timeout=180):
    r = subprocess.run(DRIVE + list(args), cwd=CWD, capture_output=True,
                       text=True, timeout=timeout)
    return (r.stdout or '') + (r.stderr or '')


def parse_ifchain(pc):
    """Some TAs (CSU, OC, WSPE) dispatch with an if/else-if chain rather than a
    switch, which has no `case N:` lines at all. Match `== N` guards followed within
    a few lines by a call."""
    out = {}
    lines = pc.splitlines()
    for n, line in enumerate(lines):
        m = re.search(r'==\s*(0x[0-9A-Fa-f]+|\d+)\s*\)', line)
        if not m or not re.search(r'\bif\b', line):
            continue
        try:
            cid = int(m.group(1), 0)
        except ValueError:
            continue
        if cid > 0x2000:
            continue
        for k in range(n + 1, min(n + 4, len(lines))):
            c = re.search(r'=\s*([A-Za-z_]\w*)\s*\(', lines[k])
            if c and c.group(1) not in ('if', 'while', 'for'):
                out.setdefault(cid, c.group(1))
                break
    return out


def parse_cases(pc):
    """case <id>: ... <something> = sub_XXXX(  ->  {id: handler}"""
    out, pending = {}, []
    for line in pc.splitlines():
        m = re.search(r'case\s+(0x[0-9A-Fa-f]+|\d+)[uU]?:', line)
        if m:
            pending.append(int(m.group(1), 0))
            continue
        m2 = (re.search(r'=\s*([A-Za-z_]\w*)\s*\(', line)
              or re.search(r'\)\s*([A-Za-z_]\w*)\s*\)\s*\(', line))
        if m2 and pending:
            for cid in pending:
                out.setdefault(cid, m2.group(1))
            pending = []
            continue
        m3 = re.search(r'return\s+([A-Za-z_]\w*)\s*\(', line)
        if m3 is None and 'return' in line:
            m3 = re.search(r'\)\s*([A-Za-z_]\w*)\s*\)\s*\(', line)
        if m3 and pending:
            for cid in pending:
                out.setdefault(cid, m3.group(1))
            pending = []
        if re.match(r'\s*(default|LABEL_|\})', line):
            pending = []
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix', required=True)
    ap.add_argument('--dispatch', required=True)
    ap.add_argument('--note')
    ap.add_argument('--out')
    a = ap.parse_args()

    def promote_locs(text):
        """Handlers are often reached as `loc_XXXX)(args)` because IDA never made a
        function there; the decompiler then shows a cast-through-label instead of a
        callee, and no rename is possible. Promote them to functions first."""
        locs = set(re.findall(r'loc_([0-9A-Fa-f]+)\s*\)\s*\(', text))
        for L in locs:
            drive('raw', 'define', 'kind=func', 'target=0x%s' % L)
        return len(locs)

    pc = drive('pc', a.dispatch)
    n_promoted = promote_locs(pc)
    if n_promoted:
        pc = drive('pc', a.dispatch)
    cases = parse_cases(pc) or parse_ifchain(pc)
    disp = a.dispatch
    # Wrapper-following: several TAs' invoke entry only validates paramTypes and then
    # tail-calls (b.w) or calls the real dispatcher. Follow one level when the entry
    # itself shows no cases.
    if not cases:
        inner = re.findall(r'(?:return|=)\s*([A-Za-z_]\w*)\s*\(', pc)
        inner = [x for x in dict.fromkeys(inner)
                 if x not in ('if', 'while', 'for', 'return', 'sizeof')]
        for cand in inner:
            pc2 = drive('pc', cand)
            if promote_locs(pc2):
                pc2 = drive('pc', cand)
            c2 = parse_cases(pc2) or parse_ifchain(pc2)
            if c2:
                if re.fullmatch(r'sub_[0-9A-Fa-f]+', cand):
                    import tempfile as _tf
                    with _tf.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
                        json.dump([dict(addr=int(cand.split('_')[1], 16),
                                        name='%s_dispatch' % a.prefix)], fh)
                        dp = fh.name
                    drive('syms', dp)
                    os.unlink(dp)
                    cand = '%s_dispatch' % a.prefix
                cases, pc, disp = c2, pc2, cand
                break
    # a handler shared by several ids gets the lowest id in its name
    byfn = {}
    for cid, fn in sorted(cases.items()):
        byfn.setdefault(fn, cid)
    # Rename by ADDRESS via rename_many, not by token: `drive mv` resolves the old
    # name against the *current view*, so it silently fails for anything off-screen
    # (8 of 14 handlers on the first GCP run).
    # Resolve EVERY handler to an address so the decompiler-derived map is
    # authoritative. The earlier static pass guessed names from a jump-table walk and
    # got some wrong (it labelled OPM's precheck `opm_cmd00`); those must be corrected,
    # not skipped just because the function already has a name.
    name2addr = {}
    for ln in drive('names', a.prefix, '400').splitlines():
        mm = re.match(r'\s*(0x[0-9A-Fa-f]+)\s+(\S+)', ln)
        if mm:
            name2addr[mm.group(2)] = int(mm.group(1), 16)
    syms, already = [], []
    for fn, cid in byfn.items():
        if re.fullmatch(r'sub_[0-9A-Fa-f]+', fn):
            addr = int(fn.split('_')[1], 16)
        elif fn in name2addr:
            addr = name2addr[fn]
            already.append(fn)
        else:
            continue
        syms.append(dict(addr=addr, name='%s_cmd%02X' % (a.prefix, cid)))
    import tempfile
    with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
        json.dump(syms, fh)
        sp = fh.name
    res = drive('syms', sp)
    os.unlink(sp)
    m = re.search(r'(\d+)/(\d+) renamed', res)
    ok = int(m.group(1)) if m else 0
    if n_promoted:
        print('   promoted %d loc_ labels to functions' % n_promoted)
    print('%s: %d commands, %d distinct handlers, %d renamed%s'
          % (a.prefix, len(cases), len(byfn), max(ok, 0),
             (', %d already named' % len(already)) if already else ''))
    print('   ids: %s' % ', '.join(hex(c) for c in sorted(cases)))
    if a.note:
        drive('note', disp, a.note)
    drive('save')
    if a.out:
        json.dump({hex(k): v for k, v in sorted(cases.items())},
                  open(a.out, 'w'), indent=1)
    return 0


if __name__ == '__main__':
    raise SystemExit(main())