#!/usr/bin/env python3
"""Find handlers where a REJECTED request still falls through to state mutation.

Generalises the GCP FE0D finding (`fw/FINDING-GCP-FE0D-REFCOUNT-DESYNC.md`). There, a
handle check was compiled as an `ITTE` block that merely *selects an error code*, so the
failure path continued into the refcount decrement and the ASD release call. Its sibling
FE0B branched away instead and was safe. The difference is visible in the instruction
stream:

    safe:    CMP ...; BNE error_label          (failure leaves the mutating code)
    unsafe:  CMP ...; IT{T,E}; MOVcc Rx,#err   (failure only sets a value, then falls through)

So the detector is: inside a command handler, an `IT`-block that materialises an
error-looking constant, followed - before any return - by a store to a **fixed context
offset** (`STR Rn,[Rm,#imm]` or a pc-relative global), i.e. persistent state changed after
a rejection.

It is a heuristic and reports leads, not verdicts: conditional execution around an error
code is perfectly normal when nothing is mutated afterwards. Every hit needs the same
sibling-comparison treatment FE0D got. Reviewed hits are recorded in the accompanying doc.

Function boundaries come from IDA (`fw/libdumps/*.json`), so decoding is correctly aligned.

usage: ta_fallthrough_audit.py [--prefix X] [--min-err 0x20]
"""
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)
IT = re.compile(r'^it[te]*$')
# capstone writes the condition BEFORE the width suffix ('movne.w', not 'mov.wne'),
# which an earlier pattern got backwards and so found nothing -- it failed its own
# control (it must rediscover GCP FE0D) and that is how the bug was caught.
MOVCC = re.compile(r'^mov(eq|ne|cs|cc|hi|ls|ge|lt|gt|le)(?:s|\.w)?$')
STORE = re.compile(r'^str[bh]?(?:\.w)?$')


def audit_fn(body, ea, size):
    """Return list of (it_addr, err_value, store_addr, store_text)."""
    ins = list(MD.disasm(body[ea - BASE:ea - BASE + size], ea))
    hits = []
    for n, i in enumerate(ins):
        if not IT.match(i.mnemonic):
            continue
        # error constant materialised under predication?
        err = None
        for k in range(n + 1, min(n + 5, len(ins))):
            m = MOVCC.match(ins[k].mnemonic)
            if m and '#' in ins[k].op_str:
                try:
                    v = int(ins[k].op_str.split('#')[1], 0)
                except ValueError:
                    continue
                if v >= 0x20:            # error codes here are 0x100, 0x2xx, 1043...
                    err = v
        if err is None:
            continue
        # after the IT block, is persistent state stored before a return?
        for k in range(n + 1, len(ins)):
            t = ins[k]
            if t.mnemonic.startswith('pop') and 'pc' in t.op_str:
                break
            if t.mnemonic == 'bx' and 'lr' in t.op_str:
                break
            # skip predicated stores: those belong to the success arm of the IT block
            if (STORE.match(t.mnemonic) and '[' in t.op_str
                    and 'sp' not in t.op_str):
                hits.append((i.address, err, t.address,
                             '%s %s' % (t.mnemonic, t.op_str)))
                break
    return hits


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix')
    ap.add_argument('--all', action='store_true',
                    help='every non-library function, not just command handlers')
    a = ap.parse_args()
    dumps = sorted(glob.glob(os.path.join(REPO, 'fw', 'libdumps', '*.json')))
    total = 0
    for p in dumps:
        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()
        rows = []
        for f in fns:
            # only handlers and their attributed helpers -- not library code
            nm = f['name']
            if 'lib_' in nm:
                continue            # shared runtime, not TA request handling
            if not a.all and '_cmd' not in nm and '_dispatch' not in nm:
                continue
            for it, err, st, txt in audit_fn(body, f['ea'], f['size']):
                rows.append((nm, it, err, st, txt))
        if rows:
            print('=== %s' % ta)
            for nm, it, err, st, txt in rows:
                print('   %-46s IT@%#x err=%#-6x -> store@%#x  %s'
                      % (nm[:46], it, err, st, txt))
            total += len(rows)
    print('\n%d lead(s). Each needs a sibling/branch comparison before it is a finding.'
          % total)
    return 0


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