#!/usr/bin/env python3
"""Derive an IDA symbol set for a PSP TA body: entry points + per-command handlers.

Feeds `idatui.drive syms <file.json>` (RPC `rename_many`), so a whole TA can be
annotated from static analysis instead of hand-navigation.

What it recovers, all from the layout established on HDCP in idatui:
  * the five TEE entry points, via the **ARM-mode** trampoline (the rest of the image is
    Thumb; decoding the trampoline as Thumb yields convincing garbage);
  * the invoke dispatcher;
  * the command -> handler map, from either dispatch form:
      - `tbb`/`tbh` jump table, with the index base taken from the preceding
        `sub rX, rY, #base` and the arm count from the `cmp rX, #n` bound;
      - a `cmp #id` / `beq` chain.
    Each arm/target is then scanned forward for the `bl` that does the work.

Ground truth check (`--selftest`): on HDCP 23.0.4.49 this must recover the dispatcher at
0x4013b4 and the command set {1..7,9,0xB,0xC,0xE,0xF,0x10,0x11,0x12,0x13} that was read
by hand out of the decompiler.

usage: ta_symbols.py <body.bin> --name <prefix> [-o syms.json]
       ta_symbols.py --selftest
"""
import argparse
import json
import os
import re
import struct
import sys

import capstone as cs

BASE = 0x400100
MT = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
MA = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)


def sweep(b, start_va, end_va):
    out, off, end = [], start_va - BASE, end_va - BASE
    while off < min(end, len(b) - 2):
        got = None
        for i in MT.disasm(b[off:off + 4], BASE + off):
            got = i
            break
        if got:
            out.append(got)
            off += got.size
        else:
            off += 2
    return out


def header(b):
    g = lambda o: struct.unpack_from('<I', b, o)[0]
    return dict(base=g(0x04), entry=g(0x10), magic=b[0x14:0x18],
                session_tbl=g(0x1c), guard=g(0x20))


def entry_points(b, h):
    ins = list(MA.disasm(b[h['entry'] - BASE:h['entry'] - BASE + 0xC0], h['entry']))
    seen, out = set(), []
    for i in ins:
        if i.mnemonic.startswith('ldr') and '[pc' in i.op_str:
            m = re.search(r'#(-?0x[0-9a-f]+|-?\d+)', i.op_str.split('[pc')[1])
            if not m:
                continue
            lit = ((i.address + 8) & ~3) + int(m.group(1), 0)
            if 0 <= lit - BASE < len(b) - 4:
                v = struct.unpack_from('<I', b, lit - BASE)[0] & ~1
                if h['base'] < v < h['base'] + len(b) and v not in seen:
                    seen.add(v)
                    out.append(v)
    return out


def first_call(b, va, span=0x28):
    for i in sweep(b, va, va + span):
        if i.mnemonic in ('bl', 'blx') and i.op_str.startswith('#'):
            try:
                return int(i.op_str[1:], 0)
            except ValueError:
                return None
    return None


def dispatch_map(b, fn, span=0x800):
    """command id -> (arm_va, handler_va or None)."""
    ins = sweep(b, fn, fn + span)
    out = {}
    for n, i in enumerate(ins):
        if i.mnemonic in ('tbb', 'tbh'):
            half = i.mnemonic == 'tbh'
            # bound: nearest preceding cmp #n ; index base: nearest preceding sub #k
            cnt, sub = None, 0
            for k in range(n - 1, max(0, n - 12), -1):
                p = ins[k]
                if cnt is None and p.mnemonic.startswith('cmp') and '#' in p.op_str:
                    try:
                        cnt = int(p.op_str.split('#')[1], 0) + 1
                    except ValueError:
                        pass
                if p.mnemonic in ('sub', 'sub.w', 'subw', 'subs') and '#' in p.op_str:
                    try:
                        sub = int(p.op_str.split('#')[1], 0)
                        break
                    except ValueError:
                        pass
            if not cnt or cnt > 64:
                continue
            tb = i.address + 4
            for k in range(cnt):
                o = (tb + (2 * k if half else k)) - BASE
                if o + 2 > len(b):
                    break
                e = struct.unpack_from('<H', b, o)[0] if half else b[o]
                arm = tb + 2 * e
                out[sub + k] = (arm, first_call(b, arm))
        if i.mnemonic.startswith('cmp') and '#' in i.op_str and n + 1 < len(ins):
            nxt = ins[n + 1]
            if nxt.mnemonic.startswith('b') and nxt.op_str.startswith('#') and \
                    nxt.mnemonic not in ('bl', 'blx', 'bx'):
                try:
                    cid = int(i.op_str.split('#')[1], 0)
                    tgt = int(nxt.op_str[1:], 0)
                except ValueError:
                    continue
                if cid < 0x200 and cid not in out:
                    out[cid] = (tgt, first_call(b, tgt))
    return out


def build(path, prefix):
    b = open(path, 'rb').read()
    h = header(b)
    eps = entry_points(b, h)
    syms, meta = [], dict(header=h, entry_points=eps, path=path)
    labels = ['TA_CreateEntryPoint', 'TA_DestroyEntryPoint', 'TA_OpenSessionEntryPoint',
              'TA_CloseSessionEntryPoint', 'TA_InvokeCommandEntryPoint']
    syms.append(dict(addr=h['entry'], name='%s_entry_trampoline' % prefix))
    for lab, ep in zip(labels, eps):
        syms.append(dict(addr=ep, name='%s_%s' % (prefix, lab)))
    cmds = {}
    if len(eps) >= 5:
        cmds = dispatch_map(b, eps[4])
        # the dispatcher is often a thin wrapper; follow one level if it has a single call
        if not cmds:
            inner = first_call(b, eps[4], 0x40)
            if inner:
                cmds = dispatch_map(b, inner)
                if cmds:
                    syms.append(dict(addr=inner, name='%s_dispatch' % prefix))
                    meta['inner_dispatch'] = inner
        # Deliberately NOT naming command handlers here. The jump-table walk guesses
        # the handler as "first bl after the arm", which is wrong whenever the arm sets
        # up arguments or falls through, and those wrong names then persist in the IDB
        # and shadow the decompiler-derived map. fw/annotate_ta.py owns command naming.
    meta['commands'] = {hex(k): (hex(v[0]), hex(v[1]) if v[1] else None)
                        for k, v in sorted(cmds.items())}
    # de-duplicate: IDA rejects a repeated name
    seen, uniq = set(), []
    for s in syms:
        if s['addr'] in seen:
            continue
        seen.add(s['addr'])
        uniq.append(s)
    return uniq, meta


def selftest():
    p = 'psp_out/versions_win/ta_0.7.0.1_543e8a45.bin'
    if not os.path.exists(p):
        print('selftest: HDCP reference body missing')
        return 1
    syms, meta = build(p, 'hdcp')
    inv = [s for s in syms if s['name'].endswith('InvokeCommandEntryPoint')]
    want = {1, 2, 3, 4, 5, 6, 7, 9, 0xB, 0xC, 0xE, 0xF, 0x10, 0x11, 0x12, 0x13}
    got = {int(k, 16) for k in meta['commands']}
    ok_inv = inv and inv[0]['addr'] == 0x4013b4
    missing = want - got
    print('selftest HDCP 23.0.4.49:')
    print('   invoke entry %s (want 0x4013b4) -> %s'
          % (hex(inv[0]['addr']) if inv else None, 'OK' if ok_inv else 'FAIL'))
    print('   commands recovered: %s' % sorted(hex(x) for x in got))
    print('   hand-read set covered: %s' % ('OK' if not missing else
                                            'MISSING %s' % sorted(hex(x) for x in missing)))
    return 0 if (ok_inv and not missing) else 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('body', nargs='?')
    ap.add_argument('--name', default='ta')
    ap.add_argument('-o', '--out')
    ap.add_argument('--selftest', action='store_true')
    a = ap.parse_args()
    if a.selftest:
        return selftest()
    syms, meta = build(a.body, a.name)
    if a.out:
        json.dump(syms, open(a.out, 'w'), indent=1)
        json.dump(meta, open(a.out.replace('.json', '.meta.json'), 'w'),
                  indent=1, default=str)
    print(json.dumps(meta['commands'], indent=1))
    print('%d symbols' % len(syms))
    return 0


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