#!/usr/bin/env python3
"""DMCUB (Xtensa) static analysis: functions, call graph, and annotated listings.

capstone-6 decodes Xtensa mnemonics/operands, but its branch and call operands
are unreliable (`op_str` shows a displacement and `operands[].imm` is wrong), so
call/jump targets are decoded here from the raw encodings:

    CALLn    op0=5, n=bits5:4, imm18=bits23:6   target = (PC&~3) + (imm18<<2) + 4
    J        op0=6, n=0,       imm18=bits23:6   target = PC + 4 + imm18
    B*(RRI8) op0=6/7 family,   imm8=bits23:16   target = PC + 4 + imm8

Requires /home/user/.venvs/cs6/bin/python (capstone >= 6).

usage:
  dmcub_dis.py <image> funcs                   list functions with call counts
  dmcub_dis.py <image> hot [N]                 hottest call targets
  dmcub_dis.py <image> dis <va> [<va>...]      annotated listing of a function
  dmcub_dis.py <image> xrefs <va>              callers of a function
  dmcub_dis.py <image> grep <regex>            functions containing a mnemonic/op pattern
"""
import collections
import importlib.util
import os
import re
import struct
import sys

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


def w24(p, i):
    return p[i] | (p[i + 1] << 8) | (p[i + 2] << 16)


class NarrowBranch:
    """capstone 6 cannot decode Xtensa's density-option branches BEQZ.N/BNEZ.N
    (op0=0xC with r >= 8) -- ~700 of them in a DCN 3.2 image, i.e. a large slice
    of the control flow, including the `beqz.n a9` that selects
    DMUB_CMD__ABM_INIT_CONFIG. This shim decodes them.

    RI6: word = b0 | b1<<8 ; op0 = word & 0xF = 0xC
         r = (word >> 4) & 0xF     s = (word >> 8) & 0xF   imm4 = word >> 12
         imm6 = ((r & 3) << 4) | imm4        target = PC + 4 + imm6
         r in 8..B -> beqz.n, r in C..F -> bnez.n   (r < 8 is movi.n)
    """

    __slots__ = ('address', 'size', 'bytes', 'mnemonic', 'op_str', 'operands', 'target_va')

    def __init__(self, off_bytes, va):
        word = off_bytes[0] | (off_bytes[1] << 8)
        r = (word >> 4) & 0xF
        s = (word >> 8) & 0xF
        imm4 = (word >> 12) & 0xF
        imm6 = ((r & 3) << 4) | imm4
        self.address = va
        self.size = 2
        self.bytes = bytes(off_bytes[:2])
        self.mnemonic = 'beqz.n' if (r & 0xC) == 0x8 else 'bnez.n'
        self.target_va = va + 4 + imm6
        self.op_str = 'a%d, . 0x%08x' % (s, self.target_va)
        self.operands = []

    @staticmethod
    def match(p, i):
        return (p[i] & 0xF) == 0xC and ((p[i] >> 4) & 0xF) >= 8


def sext(v, bits):
    return v - (1 << bits) if v >> (bits - 1) else v


class Analysis:
    """Function set + call graph.

    The function set starts from the dense run of `entry` prologues (the main code
    region) and is then grown to a fixpoint with every call target that also has
    an `entry` prologue. The growth step matters: some images keep code *below*
    the dense run -- on dcn_3_2_1 `memcpy` itself lives there, and a single-cutoff
    heuristic dropped it along with its 51 call sites. Stray `36 41 xx` matches
    inside literal pools are not call targets, so they stay excluded.
    """

    def __init__(self, path):
        self.img = cfg.Image(path)
        p = self.img.payload
        all_ents = self.img.entries()
        self.code_lo = self.img.code_start(all_ents)
        self.ent_set = set(all_ents)
        self._dis = {}
        self._reach = {}
        ents = [e for e in all_ents if e >= self.code_lo]
        for _ in range(8):
            self._set_funcs(ents, len(p))
            self._build_calls()
            extra = set()
            for t in self.calls:
                off = t - self.img.base
                if off in self.ent_set and t not in self.funcs:
                    extra.add(off)
            if not extra:
                break
            ents = sorted(set(ents) | extra)
            self._dis = {}

    def _set_funcs(self, ents, end):
        self.ents = ents
        self.funcs = {}
        for i, s in enumerate(ents):
            e = ents[i + 1] if i + 1 < len(ents) else end
            self.funcs[self.img.va(s)] = (s, e)

    def _reset_calls(self):
        self.calls = collections.Counter()
        self.callers = collections.defaultdict(set)
        self.callsites = collections.defaultdict(list)

    # ---- decoding -------------------------------------------------------
    STOP = ('retw.n', 'retw', 'ret.n', 'ret', 'j', 'jx', 'ill')

    def _decode_one(self, off, end):
        p = self.img.payload
        if off >= end:
            return None
        ins = next(self.img.md.disasm(p[off:min(off + 3, end)], self.img.va(off), 1), None)
        if ins is not None and ins.mnemonic not in ('.byte', '.db'):
            return ins
        if off + 1 < end and NarrowBranch.match(p, off):
            return NarrowBranch(p[off:off + 2], self.img.va(off))
        return ins

    def dis(self, fva):
        """Decode a function by **recursive descent**, then linear-fill the rest.

        Three problems forced this shape, each of which silently corrupted
        earlier results:

        1. capstone's `skipdata` advances one byte on an unknown opcode, so one
           undecodable instruction desynchronises the rest of the function (41%
           of the code, most of the call graph). Xtensa length depends only on
           op0, so a failed decode advances architecturally.
        2. A *successful* decode can still be garbage that straddles a basic-block
           boundary: an alignment byte before a switch case label gets absorbed
           into a bogus 3-byte instruction, eating the real first instruction of
           the case. That is how the ABM handler's `addi.n a10, a2, 4` vanished,
           which in turn hid a known finding from the taint analysis. Branch,
           call and jump-table targets are therefore used as decode anchors.
        3. Decoding past the end of the body into data invents `jx` sites, so a
           decoded `ill` terminates a path.
        """
        if fva in self._dis:
            return self._dis[fva]
        s, e = self.funcs[fva]
        base = self.img.base
        p = self.img.payload
        found = {}
        skipped = set()
        work = [s]
        while work:
            off = work.pop()
            if off in found or off in skipped or not (s <= off < e):
                continue
            ins = self._decode_one(off, e)
            if ins is None or ins.mnemonic in ('.byte', '.db'):
                # Step over an opcode capstone does not know and keep walking:
                # stopping here truncated the command dispatcher after 20
                # instructions and hid every case stub from the taint analysis.
                skipped.add(off)
                work.append(off + (2 if 8 <= (p[off] & 0xF) <= 0xD else 3))
                continue
            found[off] = ins
            m = ins.mnemonic
            if m == 'ill':
                del found[off]
                continue
            # branch/jump targets inside this function are anchors
            t = self.target(ins)
            if t is not None and s <= t - base < e and not m.startswith('call'):
                work.append(t - base)
            if m == 'jx':
                for tbl in self._jx_table_targets(found, off, s, e):
                    work.append(tbl - base)
            if m not in ('j', 'jx', 'retw.n', 'retw', 'ret.n', 'ret'):
                work.append(off + ins.size)
        # Instructions proven reachable by the descent, kept apart from the
        # linear fill below: only these are real code. Audits that must not be
        # fooled by data (jump-table scans, taint) use dis_reachable().
        self._reach[fva] = set(found)
        # linear fill of bytes no path reached (jump-table bodies, cold code)
        i = s
        while i < e:
            if i in found:
                i += found[i].size
                continue
            ins = self._decode_one(i, e)
            if ins is not None and ins.mnemonic not in ('.byte', '.db', 'ill'):
                nxt = i + ins.size
                if not any(k in found for k in range(i, nxt)):
                    found[i] = ins
                i = nxt
            else:
                i += 2 if 8 <= (p[i] & 0xF) <= 0xD else 3
        out = [found[k] for k in sorted(found)]
        self._dis[fva] = out
        return out

    def dis_reachable(self, fva):
        """Only the instructions the recursive descent actually reached."""
        ins = self.dis(fva)
        reach = self._reach.get(fva, set())
        return [x for x in ins if (x.address - self.img.base) in reach]

    def indirect_targets(self, fva):
        """callx site VA -> candidate callees.

        Sub-command dispatch in this firmware goes through tables of function
        pointers (`l32r base; l32i fp, base[+k]; callx8 fp`), e.g. the PSR op
        table at 0x6001a620 and the Replay one at 0x6001a680. Without resolving
        these, a handler's reachable set is just its own stub and an audit of it
        is worthless.

        Resolution: nearest preceding `l32r` literal is the table base; a constant
        `l32i` displacement selects one entry, otherwise every consecutive entry
        that is a known function start is a candidate.
        """
        out = {}
        ins = self.dis_reachable(fva)
        for i, x in enumerate(ins):
            if not x.mnemonic.startswith('callx'):
                continue
            reg = x.op_str.strip()
            base = None
            disp = None
            track = {reg}
            for y in range(i - 1, max(-1, i - 30), -1):
                z = ins[y]
                p = [q.strip() for q in z.op_str.split(',')]
                if not p or p[0] not in track:
                    continue
                if z.mnemonic in ('l32i', 'l32i.n') and len(p) == 3:
                    try:
                        disp = int(p[2], 0) if disp is None else disp
                    except ValueError:
                        pass
                    track.add(p[1])
                elif z.mnemonic == 'l32r':
                    base = self.lit(z)
                    break
                elif z.mnemonic in ('mov', 'mov.n') and len(p) == 2:
                    track.add(p[1])
                elif z.mnemonic == 'addx4' and len(p) == 3:
                    track.add(p[2])
                    disp = None          # variable index -> whole table
            if base is None:
                continue
            cands = []
            if disp:
                v = self.img.word(base + disp)
                if v in self.funcs:
                    cands = [v]
            if not cands:
                for n in range(64):
                    v = self.img.word(base + 4 * n)
                    if v is None:
                        break
                    if v in self.funcs:
                        cands.append(v)
                    elif v != 0:
                        break
            if cands:
                out[x.address] = cands
        return out

    def _jx_table_targets(self, found, jx_off, s, e):
        """Entries of the jump table feeding a `jx`, so case bodies get anchored."""
        base = None
        for k in sorted(found, reverse=True):
            if k >= jx_off:
                continue
            if found[k].mnemonic == 'l32r':
                base = self.lit(found[k])
                break
        if base is None:
            return []
        out = []
        for n in range(256):
            v = self.img.word(base + 4 * n)
            if v is None or not (self.img.base <= v < self.img.base + len(self.img.payload)):
                break
            out.append(v)
        return out

    def coverage(self):
        """Decode success over the code region: bytes consumed by successfully
        decoded instructions vs bytes stepped over.

        Reported this way rather than as a fraction of function spans, because
        function spans run to the next `entry` prologue and therefore include
        inter-function data, which makes a span-based figure meaningless.
        """
        p = self.img.payload
        dec = tot = 0
        i = self.code_lo
        end = len(p)
        while i < end - 3:
            ins = self._decode_one(i, end)
            if ins is not None and ins.mnemonic not in ('.byte', '.db'):
                dec += ins.size
                tot += ins.size
                i += ins.size
            else:
                step = 2 if 8 <= (p[i] & 0xF) <= 0xD else 3
                tot += step
                i += step
        return dec, tot

    def target(self, ins):
        """Resolved branch/call target VA, or None."""
        if isinstance(ins, NarrowBranch):
            return ins.target_va
        b = ins.bytes
        if not b:
            return None
        op0 = b[0] & 0xF
        pc = ins.address
        if ins.mnemonic.startswith('call') and not ins.mnemonic.startswith('callx'):
            if op0 != 5 or len(b) < 3:
                return None
            imm18 = sext((w24(b, 0) >> 6) & 0x3FFFF, 18)
            return (pc & ~3) + (imm18 << 2) + 4
        if ins.mnemonic == 'j' or ins.mnemonic.startswith('b'):
            # Xtensa has two wide branch encodings -- RRI8 (imm8 at bits 23:16,
            # e.g. BEQI/BGEU) and BRI12 (imm12 at bits 23:12, e.g. BEQZ/BNEZ) --
            # and decoding every branch as RRI8 sent the recursive descent to
            # wrong addresses, which silently truncated big handlers (PSR: 59 of
            # ~300 instructions reached). capstone prints the displacement
            # relative to the instruction address for both forms, so use it.
            m = re.search(r'\.\s*([+-])\s*(0x[0-9a-fA-F]+|\d+)\s*$', ins.op_str)
            if m:
                disp = int(m.group(2), 0)
                return pc + (disp if m.group(1) == '+' else -disp)
            m = re.search(r'0x([0-9a-fA-F]{8})\s*$', ins.op_str)
            if m:
                return int(m.group(1), 16)
            return None
        return None

    def _build_calls(self):
        self._reset_calls()
        for fva in self.funcs:
            for x in self.dis(fva):
                if x.mnemonic.startswith('call') and not x.mnemonic.startswith('callx'):
                    t = self.target(x)
                    if t is None:
                        continue
                    self.calls[t] += 1
                    self.callers[t].add(fva)
                    self.callsites[t].append((fva, x.address))

    # ---- helpers --------------------------------------------------------
    def lit(self, ins):
        """Value loaded by an l32r, or None."""
        if ins.mnemonic != 'l32r':
            return None
        try:
            la = int(ins.op_str.split()[-1], 0)
        except ValueError:
            return None
        return self.img.word(la)

    def func_of(self, va):
        best = None
        for f, (s, e) in self.funcs.items():
            if f <= va < self.img.va(e):
                if best is None or f > best:
                    best = f
        return best

    def listing(self, fva, annotate=True):
        out = []
        for x in self.dis(fva):
            note = ''
            t = self.target(x)
            if t is not None:
                note = ' -> 0x%08x%s' % (t, '  [fn]' if t in self.funcs else '')
            v = self.lit(x)
            if v is not None:
                note = ' = 0x%08x' % v
            out.append('%08x  %-14s %-9s %-28s%s' % (
                x.address, x.bytes.hex(), x.mnemonic, x.op_str, note if annotate else ''))
        return out


def main():
    path, cmd = sys.argv[1], sys.argv[2]
    a = Analysis(path)
    if cmd == 'funcs':
        print('%d functions, code region from payload 0x%x' % (len(a.funcs), a.code_lo))
        for f in sorted(a.funcs):
            s, e = a.funcs[f]
            print('  0x%08x size=%-6d calls_in=%-4d' % (f, e - s, a.calls[f]))
    elif cmd == 'hot':
        n = int(sys.argv[3]) if len(sys.argv) > 3 else 20
        for t, c in a.calls.most_common(n):
            s = a.funcs.get(t)
            print('  0x%08x calls=%-5d callers=%-4d size=%s' % (
                t, c, len(a.callers[t]), (s[1] - s[0]) if s else 'not-a-known-entry'))
    elif cmd == 'dis':
        for v in sys.argv[3:]:
            fva = int(v, 0)
            if fva not in a.funcs:
                fva2 = a.func_of(fva)
                print('# 0x%x is inside 0x%08x' % (fva, fva2))
                fva = fva2
            s, e = a.funcs[fva]
            print('=== 0x%08x  size=%d  callers=%d' % (fva, e - s, len(a.callers[fva])))
            print('\n'.join(a.listing(fva)))
    elif cmd == 'xrefs':
        t = int(sys.argv[3], 0)
        for f, site in sorted(a.callsites[t]):
            print('  0x%08x called from 0x%08x (fn 0x%08x)' % (t, site, f))
    elif cmd == 'grep':
        pat = re.compile(sys.argv[3])
        for f in sorted(a.funcs):
            for x in a.dis(f):
                line = '%s %s' % (x.mnemonic, x.op_str)
                if pat.search(line):
                    print('  fn 0x%08x @0x%08x  %s' % (f, x.address, line))
    else:
        print(__doc__)


if __name__ == '__main__':
    main()