#!/usr/bin/env python3
"""Per-command copy audit + dispatcher annotation for a TA open in idatui.

For each `<prefix>_cmdNN` handler this reads the decompiled body and flags the pattern
that produced every memory-corruption finding in this project so far: a copy whose
**length is not a literal**, i.e. derived from caller-supplied data.

It also records, as an IDA function comment on each flagged handler, the exact call and
the guard lines seen in that function, so the annotation carries the reasoning and not
just a verdict.

Heuristic, and deliberately reported as such: a non-literal length is a *lead*, not a
bug. Guards may bound it elsewhere (HDCP has both bounded `ubfx` and unbounded `ldrb`
paths in the same function), so each hit needs reading. The value is that it reduces
hundreds of functions to a short list.

usage: ta_cmd_audit.py --prefix xgmi [--note-dispatch NAME]
"""
import argparse
import json
import os
import re
import subprocess
import sys

CWD = os.path.expanduser('~/dev/ida-tui-maybe')

# The copy helper in these TAs is an unnamed thunk (HDCP's memcpy is the 4-byte
# sub_401058), so keying on the name finds nothing. Match ANY call and filter on shape:
# >= 3 arguments with a non-literal last argument is the length-bearing pattern.
COPY_RE = re.compile(r'\b([A-Za-z_]\w*)\s*\(([^;()]*(?:\([^;()]*\)[^;()]*)*)\)')
SKIP = {'if', 'while', 'for', 'return', 'switch', 'sizeof', '__fastcall', 'memset'}
LITERAL = re.compile(r'^\s*(0x[0-9A-Fa-f]+|\d+)\s*$')


def drive(*args, timeout=240):
    r = subprocess.run([sys.executable, '-m', 'idatui.drive'] + list(args),
                       cwd=CWD, capture_output=True, text=True, timeout=timeout)
    return (r.stdout or '') + (r.stderr or '')


def audit_fn(pc):
    """Return [(callee, args, lenarg)] for copies with a non-literal length."""
    hits = []
    for m in COPY_RE.finditer(pc):
        callee, args = m.group(1), m.group(2)
        if callee in SKIP:
            continue
        parts = [a.strip() for a in re.split(r',(?![^()]*\))', args)]
        if len(parts) < 3:
            continue
        ln = parts[-1]
        ln = re.sub(r'^a\d+:\s*', '', ln)
        if not LITERAL.match(ln):
            hits.append((callee, args.strip()[:90], ln[:40]))
    return hits


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix', required=True)
    ap.add_argument('--note-dispatch')
    ap.add_argument('--out')
    a = ap.parse_args()

    names = re.findall(r'(0x[0-9A-Fa-f]+)\s+(%s_\w+)' % a.prefix,
                       drive('names', a.prefix, '400'))
    cmds = [(ad, nm) for ad, nm in names if re.search(r'_cmd[0-9A-Fa-f]+', nm)]
    report = {}
    for ad, nm in sorted(cmds, key=lambda x: x[1]):
        pc = drive('pc', nm)
        hits = audit_fn(pc)
        guards = [l.strip() for l in pc.splitlines()
                  if re.search(r'\b(if|while)\b.*(<|>|==|!=)', l)][:6]
        report[nm] = dict(addr=ad, copies=hits, guards=guards,
                          lines=len(pc.splitlines()))
        if hits:
            txt = ('COPY AUDIT: %d call(s) with a non-literal length.\\n' % len(hits))
            for c, args, ln in hits:
                txt += '  %s(%s)  len=%s\\n' % (c, args, ln)
            txt += '\\nGuards seen in this function:\\n'
            for g in guards:
                txt += '  %s\\n' % g.replace('"', "'")
            txt += ('\\nHeuristic lead, not a verdict: a non-literal length may still be '
                    'bounded elsewhere (HDCP carries bounded ubfx and unbounded ldrb '
                    'paths in one function).')
            drive('note', nm, txt)
    flagged = {k: v for k, v in report.items() if v['copies']}
    print('%s: %d command handlers, %d with non-literal-length copies'
          % (a.prefix, len(cmds), len(flagged)))
    for k, v in flagged.items():
        print('   %-42s %s' % (k, '; '.join('%s len=%s' % (c, l)
                                            for c, _, l in v['copies'])[:110]))
    if a.note_dispatch:
        drive('note', a.note_dispatch,
              'Command handlers audited for non-literal-length copies: %d of %d flagged.'
              % (len(flagged), len(cmds)))
    drive('save')
    if a.out:
        json.dump(report, open(a.out, 'w'), indent=1)
    return 0


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