#!/usr/bin/env python3
"""First-pass structural triage of a PSP TA body: entry points and command surface.

Every TA in this family shares a layout, established on HDCP 23.0.4.49 in idatui and
re-checked here against the four TAs this project already knows:

  info header at body+0 (maps to 0x400100)
      +0x04 image base   +0x10 entry trampoline   +0x14 'TA'/'DR' class magic
      +0x1c session table   +0x20 &__stack_chk_guard

  the trampoline dispatches on request[0] through a pointer table:
      1 CreateEntryPoint   2 DestroyEntryPoint   3 OpenSessionEntryPoint
      4 CloseSessionEntryPoint   5 InvokeCommandEntryPoint  (4 args)

`InvokeCommandEntryPoint` is the host-facing command surface. Its dispatcher is either a
`cmp`-chain or a `tbb`/`tbh` jump table; both are enumerated here.

usage: ta_triage.py <body.bin> [...]        (add --commands for the command dump)
"""
import argparse
import os
import re
import struct
import sys

import capstone as cs

BASE = 0x400100
MD = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)


MD_ARM = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM)


def sweep_arm(b, start_va, n=0x40):
    """The entry trampoline is ARM, not Thumb (4-byte encodings), unlike the
    rest of the image. Decoding it as Thumb yields convincing garbage."""
    off = start_va - BASE
    return list(MD_ARM.disasm(b[off:off + n * 4], start_va))


def sweep(b, start_va, end_va=None):
    """Linear sweep that resumes past undecodable words (data in .text)."""
    out = []
    off = start_va - BASE
    end = (end_va - BASE) if end_va else len(b) - 2
    while off < end:
        got = None
        for i in MD.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):
    """Walk the trampoline, collecting the literal pointers it jumps through."""
    ins = sweep_arm(b, h['entry'], 0x30)
    ptrs = []
    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 + 4) & ~3) + int(m.group(1), 0)
            if 0 <= lit - BASE < len(b) - 4:
                v = struct.unpack_from('<I', b, lit - BASE)[0]
                if h['base'] < (v & ~1) < h['base'] + len(b):
                    ptrs.append(v & ~1)
    # de-dup preserving order; the trampoline emits them in case order 1..5
    seen, out = set(), []
    for p in ptrs:
        if p not in seen:
            seen.add(p)
            out.append(p)
    return out


def commands(b, fn, limit=0x600):
    """Enumerate the command ids a dispatcher tests."""
    ins = sweep(b, fn, fn + limit)
    ids, tables = set(), []
    for n, i in enumerate(ins):
        if i.mnemonic in ('cmp', 'cmp.w') and '#' in i.op_str:
            try:
                ids.add(int(i.op_str.split('#')[1], 0))
            except ValueError:
                pass
        if i.mnemonic in ('subs', 'sub.w', 'subw') and '#' in i.op_str:
            pass
        if i.mnemonic in ('tbb', 'tbh'):
            tables.append(i.address)
    return sorted(ids), tables


def tb_entries(b, tb_va, is_half, n):
    base = tb_va + 4
    out = []
    for k in range(n):
        o = (base + (2 * k if is_half else k)) - BASE
        if o + 2 > len(b):
            break
        e = struct.unpack_from('<H', b, o)[0] if is_half else b[o]
        out.append(base + 2 * e)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('files', nargs='+')
    ap.add_argument('--commands', action='store_true')
    args = ap.parse_args()
    for p in args.files:
        b = open(p, 'rb').read()
        h = header(b)
        eps = entry_points(b, h)
        name = os.path.basename(p)
        print('\n=== %s  (%#x bytes)' % (name, len(b)))
        print('    base=%#x  entry=%#x  class=%s  session_tbl=%#x  guard=%#x'
              % (h['base'], h['entry'], h['magic'][:2].decode('latin1'),
                 h['session_tbl'], h['guard']))
        labels = ['Create', 'Destroy', 'OpenSession', 'CloseSession', 'Invoke']
        for lab, ep in zip(labels, eps):
            print('    %-13s %#x' % (lab, ep))
        if len(eps) >= 5 and args.commands:
            ids, tbs = commands(b, eps[4])
            print('    invoke dispatcher @%#x' % eps[4])
            print('      cmp constants: %s' % [hex(x) for x in ids if x < 0x200])
            for t in tbs:
                print('      jump table @%#x' % t)
    return 0


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