#!/usr/bin/env python3
"""Write the per-TA summary into the IDB as function comments.

Closes the annotation loop: everything recovered about a TA (entry points, param gates,
command map, RPC reach) is written onto the invoke entry so the database is
self-describing, rather than living only in markdown.

usage: ta_final_note.py --prefix opm   (IDATUI_RPC_SOCK -> the pane)
"""
import argparse
import json
import os
import re
import sys

CWD = os.path.expanduser('~/dev/ida-tui-maybe')
sys.path.insert(0, CWD)
from idatui.rpcclient import RpcClient                                   # noqa: E402

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


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

    cmds = {}
    cp = os.path.join(REPO, 'fw', 'cmdmaps', '%s.json' % p)
    if os.path.exists(cp):
        cmds = json.load(open(cp))
    rpc = json.load(open(os.path.join(REPO, 'fw', 'audits', 'rpc_table.json')))
    meta = json.load(open(os.path.join(REPO, 'fw', 'syms', '%s.meta.json' % p)))
    h = meta['header']

    cl = RpcClient(os.environ['IDATUI_RPC_SOCK'], timeout=60)
    cl.connect()

    fs = cl.call('functions', filter='%s_TA_InvokeCommandEntryPoint' % p, limit=4)
    if not fs:
        print('%s: no invoke entry named' % p)
        return 1
    ea = fs[0]['ea']
    d = cl.call('disassembly', target=hex(ea))
    d = d.get('disassembly', d)
    dis = [l.get('text', '') for l in d.get('lines', [])]
    gate = next((t for t in dis if re.search(r'CMP.*#0x[0-9A-Fa-f]{3,}', t)), None)

    ids = sorted(cmds, key=lambda x: int(x, 16))
    lines = [
        '%s TA -- recovered structure.' % p.upper(),
        '',
        'info header @0x400100: base=%#x entry=%#x session_tbl=%#x guard=%#x'
        % (h['base'], h['entry'], h['session_tbl'], h['guard']),
        'class magic +0x14 = %r (TA = host-invokable; ASD is DR and is not)'
        % str(meta['header']['magic'])[:12],
        '',
        'Entry points (via the ARM-mode trampoline at %#x):' % h['entry'],
    ]
    labels = ['Create', 'Destroy', 'OpenSession', 'CloseSession', 'Invoke']
    for lab, epa in zip(labels, meta.get('entry_points', [])):
        lines.append('  %-12s %#x' % (lab, epa))
    lines += ['', 'Command surface: %d ids' % len(ids)]
    for k in ids:
        lines.append('  cmd %-6s -> %s' % (k, cmds[k]))
    if gate:
        lines += ['', 'Param-size gate seen in the entry: %s' % gate.strip()]
    t = rpc.get(p, {})
    lines += ['', 'Service RPC reach (IDA-derived, fw/ta_rpc_names.py):']
    lines.append('  SYS_DRV (svc 1): %s'
                 % (', '.join(hex(x) for x in t.get('sysdrv', [])) or 'none resolved'))
    lines.append('  ASD     (svc 2): %s'
                 % (', '.join(hex(x) for x in t.get('asd', [])) or 'none resolved'))
    if 0x202 in t.get('asd', []):
        lines += ['', '!! reaches ASD 0x202 = uninitialised-stack canary disclosure,',
                  '   previously rated not-host-reachable on a four-TA scan.']
    if 0x301 in t.get('asd', []):
        lines += ['', '!! reaches ASD 0x301 = content-controlled write up to 5120 bytes,',
                  '   previously dismissed because no known TA emitted it.']

    txt = '\\n'.join(lines).replace('"', "'")
    # `note` is a drive-level convenience, not an RPC verb: it does goto(fn) then
    # comment(text) on the function's first line. Replicate that here.
    cl.call('goto', target=hex(ea), delay_ms=0)
    cl.call('comment', text=txt)
    cl.call('save')
    print('%s: annotated invoke entry @%#x with %d lines' % (p, ea, len(lines)))
    return 0


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