#!/usr/bin/env python3
"""Find slot lookups that never test whether the slot is live.

Generalises GCP FE0F (`fw/FINDING-GCP-FE0F-STALE-SLOT-LOOKUP.md`): it walks a 64-entry
array comparing the caller's bytes against a slot field and returns the slot's handle,
without ever reading the slot's state word. Combined with a destroy path that clears only
that state word, destroyed objects stay findable.

Shape of the detector, per function:
  * a **slot scan**: a loop whose back-edge is preceded by a constant stride added to a
    pointer register (`add rX, #N` / `add.w rX, rX, #N`, 8 <= N <= 256), or an index
    multiply by such a stride;
  * a **liveness test**: a `cmp` of a loaded value against a small constant (1..8) -- the
    state words in these TAs are 1 (allocated) or 2 (ready);
  * report scans that have **no** liveness test in the same function.

This is a lead generator. A scan may validate liveness in a callee, or the field it matches
on may itself be cleared on destroy (which is the correct fix and makes the scan safe). Each
hit needs the destroy path checked, exactly as FE0F needed FE0D.

usage: ta_liveness_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)
STRIDE_ADD = re.compile(r'^(add|adds|add\.w)$')
BRANCH = ('b', 'bne', 'beq', 'bhi', 'bls', 'bcs', 'bcc', 'blo', 'bge', 'blt',
          'b.w', 'bne.w', 'beq.w', 'blo.w', 'bhi.w')


def strides_and_state(ins):
    strides, states, back = [], [], False
    for n, i in enumerate(ins):
        if STRIDE_ADD.match(i.mnemonic) and '#' in i.op_str and 'sp' not in i.op_str:
            try:
                v = int(i.op_str.split('#')[1], 0)
            except ValueError:
                continue
            if 8 <= v <= 256:
                strides.append((i.address, v))
        if i.mnemonic in ('mul', 'muls', 'mla', 'mlane') or \
           (i.mnemonic.startswith('add') and 'lsl' in i.op_str):
            strides.append((i.address, -1))          # index arithmetic
        if i.mnemonic.startswith('cmp') and '#' in i.op_str:
            try:
                v = int(i.op_str.split('#')[1], 0)
            except ValueError:
                continue
            if 1 <= v <= 8:
                states.append((i.address, v))
        if i.mnemonic in BRANCH and i.op_str.startswith('#'):
            try:
                if int(i.op_str[1:], 0) < i.address:
                    back = True
            except ValueError:
                pass
    return strides, states, back


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix')
    a = ap.parse_args()
    total = 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()
        rows = []
        for f in fns:
            if 'lib_' in f['name']:
                continue
            ins = list(MD.disasm(body[f['ea'] - BASE:f['ea'] - BASE + f['size']], f['ea']))
            if not ins:
                continue
            strides, states, back = strides_and_state(ins)
            if back and strides and not states:
                rows.append((f['name'], f['ea'],
                             sorted({v for _, v in strides if v > 0})))
        if rows:
            print('=== %s' % ta)
            for nm, ea, sv in rows:
                print('   %-44s @%#x  strides=%s  NO liveness test'
                      % (nm[:44], ea, sv))
            total += len(rows)
    print('\n%d scan-without-liveness lead(s). Each needs its destroy path checked.'
          % total)
    return 0


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