#!/usr/bin/env python3
"""What can each TA ask the PSP kernel / SYS_DRV to do? (`SVC` surface per TA)

This is the analysis that decided reachability for ASD: an 89-site scan of `SVC 0xF2`
call sites across the four known TAs showed every service opcode is an immediate and no
TA forwards a host-controlled opcode. The twelve newly-hittable TAs have never been
scanned, and they are host-invokable, so their privilege reach matters.

For each body:
  * every `svc #N` site, grouped by N (0xF2 = service RPC, others = kernel calls);
  * for `svc 0xF2`, the service id (r0) and opcode, recovered by walking backwards for
    the immediate that last wrote r0/r1 — and **flagging any site where the opcode is
    not an immediate**, which is the dangerous shape (a TA forwarding a caller-chosen
    opcode into SYS_DRV).

usage: ta_svc_scan.py <body.bin> [...]
"""
import argparse
import collections
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)


def sweep(b, start=0x401000):
    out, off = [], start - BASE
    while off < 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 last_imm(ins, n, reg, depth=18):
    """Nearest preceding immediate write to `reg`; None if it came from memory/register."""
    for k in range(n - 1, max(0, n - depth), -1):
        i = ins[k]
        dst = i.op_str.split(',')[0].strip()
        if dst != reg:
            continue
        if i.mnemonic in ('mov', 'movs', 'movw', 'mov.w') and '#' in i.op_str:
            try:
                return int(i.op_str.split('#')[1], 0)
            except ValueError:
                return None
        return None          # written from something that is not an immediate
    return None


def stored_opcode(ins, n, depth=26):
    """The service opcode is not in a register: GUARD-INIT-CHAIN.md shows the ABI is
    `req[0] = opcode; r0 = service id; r1 = &req; svc 0xF2`. So recover the immediate
    most recently stored at offset 0 of some buffer before the call."""
    imm = {}
    for k in range(max(0, n - depth), n):
        i = ins[k]
        parts = [x.strip() for x in i.op_str.split(',')]
        if i.mnemonic in ('mov', 'movs', 'movw', 'mov.w') and '#' in i.op_str:
            try:
                imm[parts[0]] = int(i.op_str.split('#')[1], 0)
            except ValueError:
                pass
        if i.mnemonic.startswith('str') and len(parts) >= 2:
            mem = ','.join(parts[1:])
            if re.fullmatch(r'\[\w+\]', mem) or re.fullmatch(r'\[\w+, #0\]', mem):
                if parts[0] in imm:
                    return imm[parts[0]]
    return None


def scan(path):
    b = open(path, 'rb').read()
    ins = sweep(b)
    svcs = collections.Counter()
    rpc = collections.Counter()
    nonimm = []
    for n, i in enumerate(ins):
        if i.mnemonic != 'svc':
            continue
        try:
            num = int(i.op_str.strip('#'), 0)
        except ValueError:
            continue
        svcs[num] += 1
        if num == 0xF2:
            svc_id = last_imm(ins, n, 'r0')
            op = stored_opcode(ins, n)
            if op is None:
                nonimm.append(i.address)
            rpc[(svc_id, op)] += 1
    return svcs, rpc, nonimm, len(ins)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('bodies', nargs='+')
    a = ap.parse_args()
    for p in a.bodies:
        svcs, rpc, nonimm, n = scan(p)
        print('\n=== %s  (%d insns)' % (os.path.basename(p), n))
        print('    svc numbers: %s'
              % ', '.join('%#x x%d' % (k, v) for k, v in sorted(svcs.items())))
        if rpc:
            print('    svc 0xF2 sites: %d' % sum(rpc.values()))
            for (sid, op), c in sorted(rpc.items(), key=lambda kv: (str(kv[0][0]), str(kv[0][1]))):
                print('       service=%-6s opcode=%-8s x%d'
                      % (sid if sid is None else hex(sid),
                         op if op is None else hex(op), c))
        print('    sites whose OPCODE is not an immediate: %d %s'
              % (len(nonimm), [hex(x) for x in nonimm[:8]]))
    return 0


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