#!/usr/bin/env python3
"""Find copies whose LENGTH comes from memory (i.e. potentially from the request).

This is the shape behind every memory-corruption finding in this tree: HDCP's subcmd12
overflow is `ldrb` of a request byte, multiplied, straight into a copy length. So:

  step 1  identify each TA's copy primitives structurally -- small leaf functions whose
          body both loads and stores in a loop (`ldrb/strb` or `ldm/stm`). Naming-based
          detection does not work here: the copy helper is an unnamed thunk, and keying on
          "memcpy" in a name found nothing at all in an earlier attempt, while keying on
          "any 3-arg call" drowned in log/printf calls.
  step 2  for every call to one of those, recover how R2 (the length) was last written.
          An immediate is fine. A load from memory is a lead, and the load's addressing
          mode says where from -- `[rN,#imm]` off a request pointer is the dangerous case.

Reports leads with the instruction that produced the length, so each can be judged rather
than taken on faith. Boundaries come from IDA (`fw/libdumps/*.json`).

usage: ta_copy_audit.py [--prefix X]
"""
import argparse
import glob
import json
import os
import re

import capstone as cs

BASE = 0x400100
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MD = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)


def decode(body, ea, size):
    return list(MD.disasm(body[ea - BASE:ea - BASE + size], ea))


def is_copy_primitive(ins):
    """Loads AND stores AND a backward branch, in a small function."""
    ld = sum(1 for i in ins if i.mnemonic.startswith(('ldr', 'ldm')))
    st = sum(1 for i in ins if i.mnemonic.startswith(('str', 'stm')))
    back = any(i.mnemonic.startswith('b') and i.op_str.startswith('#') and
               int(i.op_str[1:], 0) < i.address for i in ins
               if i.op_str.startswith('#') and i.mnemonic in
               ('b', 'bne', 'beq', 'bhi', 'bls', 'bcs', 'bcc', 'b.w', 'bne.w'))
    return ld >= 1 and st >= 1 and back


def last_write(ins, n, reg, depth=14):
    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
        return i
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix')
    a = ap.parse_args()
    grand = 0
    for p in sorted(glob.glob(os.path.join(REPO, 'fw', 'libdumps', '*.json'))):
        ta = os.path.basename(p)[:-5]
        if a.prefix and ta != a.prefix:
            continue
        fns = json.load(open(p))
        meta = json.load(open(os.path.join(REPO, 'fw', 'syms', '%s.meta.json' % ta)))
        body = open(meta['path'], 'rb').read()
        byea = {f['ea']: f for f in fns}

        prims = {}
        for f in fns:
            if f['size'] > 0x80:
                continue
            ins = decode(body, f['ea'], f['size'])
            if ins and is_copy_primitive(ins):
                prims[f['ea']] = f['name']
        if not prims:
            print('=== %s: no copy primitive identified' % ta)
            continue

        leads = []
        for f in fns:
            ins = decode(body, f['ea'], f['size'])
            for n, i in enumerate(ins):
                if i.mnemonic not in ('bl', 'blx') or not i.op_str.startswith('#'):
                    continue
                try:
                    t = int(i.op_str[1:], 0)
                except ValueError:
                    continue
                if t not in prims:
                    continue
                w = last_write(ins, n, 'r2')
                if w is None:
                    leads.append((f['name'], i.address, prims[t], 'r2 unknown'))
                elif w.mnemonic.startswith('ldr'):
                    leads.append((f['name'], i.address, prims[t],
                                  '%s %s' % (w.mnemonic, w.op_str)))
        print('=== %s: %d copy primitive(s) %s ; %d length-from-memory lead(s)'
              % (ta, len(prims), [hex(x) for x in prims], len(leads)))
        for nm, at, pr, how in leads:
            print('   %-44s call@%#x -> %-18s len <- %s' % (nm[:44], at, pr[:18], how))
        grand += len(leads)
    print('\n%d lead(s) total. A load is not a bug: the value may be bounded elsewhere.'
          % grand)
    return 0


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