#!/usr/bin/env python3
"""Name a TA's service-RPC thunks. Sweep discovers candidates; **IDA decides everything.**

History, because it is the point of the design: `fw/ta_svc_scan.py` finds `svc 0xF2`
sites with a resumable linear sweep. The sites it finds are real — IDA confirmed 25/25 in
GCP — but the sweep can be **misaligned between them**, and its opcode attribution in GCP
came out shifted by one thunk, giving every thunk its neighbour's opcode. A published
table was wrong because of it.

So: the sweep contributes only "there is a service call near here". For each candidate
this asks IDA for the containing function, then reads that function's own disassembly to
recover the service id (`MOVS R0,#n` before `SVC`) and the opcode (the immediate stored
into `req[0]`), and names it `<prefix>_{asd,sysdrv}_rpc_<OPCODE>`. Anything it cannot
resolve is reset to `sub_<addr>` rather than left with a guessed name.

usage: ta_rpc_names.py --prefix gcp --body <body.bin>   (IDATUI_RPC_SOCK must point at the pane)
"""
import argparse
import json
import os
import re
import sys
import tempfile

CWD = os.path.expanduser('~/dev/ida-tui-maybe')
sys.path.insert(0, CWD)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from idatui.rpcclient import RpcClient                                   # noqa: E402

MOV_R0 = re.compile(r'MOV(?:W|S|\.W)?\s+R0,\s*#(0x[0-9A-Fa-f]+|\d+)')
SERVICE = {1: 'sysdrv', 2: 'asd'}


def connect():
    sock = os.environ.get('IDATUI_RPC_SOCK')
    if not sock:
        import glob
        c = glob.glob('/run/user/%d/idatui-*.sock' % os.getuid())
        if len(c) != 1:
            raise SystemExit('set IDATUI_RPC_SOCK (%d panes live)' % len(c))
        sock = c[0]
    cl = RpcClient(sock, timeout=60)
    cl.connect()
    return cl


FRAME = re.compile(r'SUB\s+SP,\s*SP,\s*#(0x[0-9A-Fa-f]+|\d+)')
STORE = re.compile(r'STR\s+R0,\s*\[SP,\s*#(0x[0-9A-Fa-f]+|\d+)\+var_([0-9A-Fa-f]+)\]')


def derive(cl, ea):
    """(service, opcode) from IDA's own disassembly of the function at ea.

    The opcode is req[0], i.e. the immediate stored at **[SP]** -- the store whose
    var_ offset equals the frame size. Taking the *first* stack store instead reads
    a neighbouring field: XGMI stores 2 to var_78 before storing the real opcode
    0x1067 to var_7C, and the naive rule reported every XGMI/AUC/FP opcode as 0x2.
    """
    d = cl.call('disassembly', target=hex(ea))
    d = d.get('disassembly', d)
    last = op = svc = None
    frame = None
    for ln in d.get('lines', []):
        t = ln.get('text', '')
        f = FRAME.search(t)
        if f and frame is None:
            frame = int(f.group(1), 0)
        m = MOV_R0.match(t)
        if m:
            last = int(m.group(1), 0)
        st = STORE.match(t)
        if st and op is None:
            base, var = int(st.group(1), 0), int(st.group(2), 16)
            if frame is not None and var == frame and base == frame:
                op = last
        elif t.startswith('STR R0, [SP]') and op is None:
            op = last
        if t.startswith('SVC'):
            svc = last
    return svc, op, d.get('name')


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

    from ta_svc_scan import sweep                                        # noqa: E402
    b = open(a.body, 'rb').read()
    cands = [i.address for i in sweep(b)
             if i.mnemonic == 'svc' and i.op_str.strip('#') in ('0xf2', '242')]

    cl = connect()
    confirmed, rejected, funcs = 0, [], {}
    for addr in cands:
        d = cl.call('disassembly', target=hex(addr))
        d = d.get('disassembly', d)
        txt = next((l.get('text', '') for l in d.get('lines', [])
                    if l.get('ea') == addr), '')
        if 'SVC' not in txt.upper():
            rejected.append(addr)
            continue
        confirmed += 1
        if d.get('ea') is not None:
            funcs[d['ea']] = None

    syms, table = [], {}
    for ea in sorted(funcs):
        svc, op, _ = derive(cl, ea)
        if svc in SERVICE and op is not None:
            nm = '%s_%s_rpc_%X' % (a.prefix, SERVICE[svc], op)
            table.setdefault(SERVICE[svc], []).append(op)
        else:
            nm = 'sub_%06X' % ea
        syms.append(dict(addr=ea, name=nm))
    with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
        json.dump(syms, fh)
        sp = fh.name
    r = cl.call('rename_many', file=sp) or {}
    m = r.get('rename_many', r)
    os.unlink(sp)
    cl.call('save')

    print('%s: %d svc candidates, %d confirmed by IDA, %d rejected; %d thunk functions, '
          'renamed %s/%s' % (a.prefix, len(cands), confirmed, len(rejected), len(syms),
                             m.get('ok'), m.get('requested')))
    for k in sorted(table):
        print('   %-7s opcodes: %s'
              % (k, ', '.join(hex(x) for x in sorted(set(table[k])))))
    if a.out:
        json.dump(dict(table={k: sorted(set(v)) for k, v in table.items()},
                       confirmed=confirmed, rejected=[hex(x) for x in rejected],
                       thunks=len(syms)), open(a.out, 'w'), indent=1)
    return 0


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