#!/usr/bin/env python3
"""Find every call to the DMCUB firmware's byte-copy routine and classify the
length argument, so that the host-controlled ones can be reviewed by hand.

The copy routine is located by shape, not by address: an Xtensa function with a
`loopnez` and a run of unrolled `l8ui`/`s8i` pairs (memcpy(dst=a2, src=a3, n=a4);
under call8 the caller passes dst/src/n in a10/a11/a12).

The length classifier is a backward slice inside the calling basic block:
  const    -- movi/movi.n, safe by construction
  lit      -- l32r literal
  MEMLOAD  -- l8ui/l16ui/l32i: the length came out of memory. These are the ones
              that matter: if the memory is a mailbox command frame, the length
              is host-supplied.
  <op>(..) -- derived; the detail column shows the chain

usage: dmcub_copyaudit.py <image> [--all] [--target 0xVA]
"""
import argparse
import collections
import importlib.util
import os
import re

HERE = os.path.dirname(os.path.abspath(__file__))
_s = importlib.util.spec_from_file_location('dmcub_dis', os.path.join(HERE, 'dmcub_dis.py'))
dis = importlib.util.module_from_spec(_s)
_s.loader.exec_module(dis)

ARITH = ('add', 'add.n', 'addi', 'addi.n', 'sub', 'and', 'or', 'slli', 'srli',
         'addx2', 'addx4', 'addx8', 'min', 'max', 'minu', 'maxu', 'moveqz',
         'movnez', 'extui', 'mul16s', 'mul16u', 'mull')


def ops(x):
    return [p.strip() for p in x.op_str.split(',')]


def find_memcpy(a):
    """The byte-memcpy function, by shape.

    Candidates are functions with an unrolled `l8ui`/`s8i` run plus a `loop*`.
    Several helpers match that loosely (a `loop`-based clear, inlined copies), so
    rank by *incoming call count* first and size second -- picking merely the
    smallest candidate selected a 0-caller function on dcn_3_2_1.
    """
    cands = []
    for f, (s, e) in a.funcs.items():
        ins = a.dis(f)
        pairs = sum(1 for i, x in enumerate(ins[:-1])
                    if x.mnemonic == 'l8ui' and ins[i + 1].mnemonic == 's8i')
        if pairs >= 4 and any(x.mnemonic.startswith('loop') for x in ins):
            cands.append((a.calls[f], -(e - s), f))
    return max(cands)[2] if cands else None


def classify(a, ins, idx, reg, depth=0):
    if depth > 6:
        return ('deep', '')
    for y in range(idx - 1, -1, -1):
        x = ins[y]
        p = ops(x)
        if not p or p[0] != reg:
            continue
        m = x.mnemonic
        if m in ('movi', 'movi.n'):
            try:
                return ('const', int(p[1], 0))
            except ValueError:
                return ('const?', x.op_str)
        if m == 'l32r':
            return ('lit', a.lit(x))
        if m in ('mov', 'mov.n'):
            return classify(a, ins, y, p[1], depth + 1)
        if m.startswith(('l8ui', 'l16ui', 'l32i')):
            return ('MEMLOAD', '%s %s' % (m, x.op_str))
        if m in ARITH and len(p) > 1 and re.fullmatch(r'a\d+', p[1]):
            k = classify(a, ins, y, p[1], depth + 1)
            return ('%s(%s)' % (m, k[0]), '%s | %s' % (x.op_str, k[1]))
        return (m, x.op_str)
    return ('argin/unknown', '')


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('--target', default=None, help='copy routine VA (default: detect)')
    ap.add_argument('--all', action='store_true', help='also list constant-length sites')
    args = ap.parse_args()

    a = dis.Analysis(args.image)
    tgt = int(args.target, 0) if args.target else find_memcpy(a)
    if tgt is None:
        print('no copy routine found')
        return
    sites = a.callsites.get(tgt, [])
    print('%s: copy routine 0x%08x, %d call sites' % (
        os.path.basename(args.image), tgt, len(sites)))
    rows = []
    for fva, site in sorted(sites):
        ins = a.dis(fva)
        i = [k for k, x in enumerate(ins) if x.address == site][0]
        rows.append((fva, site,
                     classify(a, ins, i, 'a12'),
                     classify(a, ins, i, 'a10'),
                     classify(a, ins, i, 'a11')))
    consts = collections.Counter(r[2][1] for r in rows if r[2][0] == 'const')
    print('constant length: %d sites %s' % (sum(consts.values()), sorted(consts)))
    print('\nnon-constant length -- review these:')
    for fva, site, n, dst, src in rows:
        if n[0] == 'const' and not args.all:
            continue
        print('  fn 0x%08x call 0x%08x' % (fva, site))
        print('      n   = %-22s %s' % (n[0], n[1]))
        print('      dst = %-22s %s' % (dst[0], dst[1]))
        print('      src = %-22s %s' % (src[0], src[1]))


if __name__ == '__main__':
    main()