#!/usr/bin/env python3
"""Audit every indirect-jump dispatch (`jx`) in a DMCUB image for a bounds check.

Only instructions reached by recursive descent are considered; a linear sweep over
inter-function data invents `jx` sites out of table bytes.

Why this matters: the firmware dispatches host-supplied command `type` and
`sub_type` bytes through jump tables (`l32r table; addx4 idx; l32i; jx`). A table
dispatch with no bound on the index is an immediate arbitrary-PC primitive from
the host mailbox, so each site must be shown to be guarded.

For each `jx` the tool recovers, by backward scan inside the function:
  * the table base (l32r literal feeding the addx4)
  * the index register
  * the table length, measured as the run of table entries that are plausible
    code addresses inside the image
  * any `bgeu/bgeui/bltu/bltui/bge/blt` on the index register (or on what it was
    copied/derived from) before the jump, and the constant it compares against

and flags a site as UNGUARDED when no such comparison is found. A flag is a
review item, not a bug: the check can live in a caller. Conversely a bound that
exceeds the table length is a real finding.

usage: dmcub_jxaudit.py <image>
"""
import importlib.util
import os
import re
import sys

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)

CMPS = ('bgeu', 'bgeui', 'bltu', 'bltui', 'bge', 'blt', 'bgei', 'blti', 'bnei', 'beqi')


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


def table_len(a, base, limit=512):
    """Entries at `base` that look like code addresses inside the image."""
    n = 0
    lo = a.img.va(a.code_lo)
    hi = a.img.va(len(a.img.payload))
    while n < limit:
        v = a.img.word(base + 4 * n)
        if v is None or not (lo <= v < hi):
            break
        n += 1
    return n


def audit(a):
    rows = []
    for fva in sorted(a.funcs):
        ins = a.dis_reachable(fva)
        for i, x in enumerate(ins):
            if x.mnemonic != 'jx':
                continue
            jreg = ops(x)[0]
            base = idx = None
            # backward: find the l32i that loaded the target, its addx4, the base
            # literal and the index register
            regs_of_interest = {jreg}
            # Walk back to the *nearest* addx4 feeding the jump; an earlier,
            # unrelated addx4 in the same function must not overwrite it (that
            # bug made the GPINT dispatch look unguarded).
            for y in range(i - 1, max(-1, i - 40), -1):
                z = ins[y]
                p = ops(z)
                if not p or p[0] not in regs_of_interest:
                    continue
                if z.mnemonic in ('l32i', 'l32i.n') and len(p) == 3:
                    regs_of_interest.add(p[1])
                elif z.mnemonic == 'addx4' and len(p) == 3 and idx is None:
                    idx = p[1]
                    regs_of_interest.add(p[2])
                elif z.mnemonic == 'l32r' and base is None:
                    base = a.lit(z)
                elif z.mnemonic in ('mov', 'mov.n') and len(p) == 2:
                    regs_of_interest.add(p[1])
            # look for a comparison on the index (or an alias) before the jump
            guard = None
            if idx:
                alias = {idx}
                for y in range(i - 1, max(-1, i - 60), -1):
                    z = ins[y]
                    p = ops(z)
                    if z.mnemonic in CMPS and len(p) >= 3 and (p[0] in alias or p[1] in alias):
                        guard = '%s %s' % (z.mnemonic, z.op_str)
                        break
                    if p and p[0] in alias and z.mnemonic in ('mov', 'mov.n', 'extui',
                                                              'addi', 'addi.n', 'and',
                                                              'l8ui', 'srli', 'slli'):
                        if len(p) > 1 and re.fullmatch(r'a\d+', p[1]):
                            alias.add(p[1])
            n = table_len(a, base) if base else 0
            rows.append((x.address, fva, base, n, idx, guard))
    return rows


def main():
    a = dis.Analysis(sys.argv[1])
    rows = audit(a)
    unguarded = 0
    print('%-11s %-11s %-11s %-6s %-5s %s' % ('jx', 'function', 'table', 'len', 'idx', 'guard'))
    for site, fva, base, n, idx, guard in rows:
        if guard is None:
            unguarded += 1
        print('0x%08x 0x%08x %-11s %-6s %-5s %s' % (
            site, fva, ('0x%08x' % base) if base else '?', n, idx or '?',
            guard or '*** none found ***'))
    print('\n%d jx sites, %d with no bound recovered' % (len(rows), unguarded))


if __name__ == '__main__':
    main()