#!/usr/bin/env python3
"""First-pass audit tooling for the MES (Micro Engine Scheduler) firmware.

MES is the RV64 scheduler microcontroller on gfx11; the driver drives it with API
packets (`mes_v11_api_def.h`), so it has the same shape of host-reachable command
surface as the DMCUB inbox that `../dmcub/` audits -- and it had never been looked
at in this tree beyond an ISA classification.

Container (verified, same nested-header trick as DMUB/MEC):
    0x000  common_firmware_header + mes_firmware_header_v1_0 (header_size 0x48)
    0x100  payload, whose first 0x100 bytes are a **PSP header** ($PS1 at +0x10)
    0x200  ucode image, mapped at `mes_uc_start_addr - 0x2000`
           (the first 0x2000 bytes are zero, so code begins at file 0x2200 =
            uc_start_addr, 0xf0005000)
    +data  data image at `mes_ucode_data_offset_bytes`, mapped at
           `mes_data_start_addr` (0xf0100000)

Requires capstone (RISC-V); the repo venv at /home/user/.venvs/cs6 works.

usage:
  mes_audit.py <image> info        container, ISA evidence, function count
  mes_audit.py <image> dispatch    every indirect jump/call with table + bound
  mes_audit.py <image> tables      code-pointer tables in the data image
"""
import argparse
import collections
import re
import struct
import sys

import capstone as cs

GUARDS = ('bltu', 'bgeu', 'blt', 'bge', 'bne', 'beq', 'bnez', 'beqz', 'bltz', 'bgez')
INDIRECT = ('jalr', 'jr', 'c.jalr', 'c.jr')


class Mes:
    """Handles both `mes_firmware_header_v1_0` (MES/MES1) and
    `gfx_firmware_header_v2_0` (PFP/ME/MEC), so the same dispatch audit runs on
    every RV64 gfx microcontroller on this ASIC."""

    def __init__(self, path):
        self.raw = open(path, 'rb').read()
        self.path = path
        (self.size, self.hdr_size, self.hv_maj, self.hv_min, self.ip_maj, self.ip_min,
         self.ucode_version, self.ucode_size, self.ucode_array_offset,
         self.crc32) = struct.unpack_from('<IIHHHHIIII', self.raw, 0)
        if self.hdr_size >= 0x48:
            # mes_firmware_header_v1_0
            self.kind = 'mes'
            self.mes_ucode_version, = struct.unpack_from('<I', self.raw, 0x20)
            (self.mes_ucode_size, self.mes_ucode_offset,
             self.mes_data_version, self.mes_data_size, self.mes_data_offset,
             self.uc_start_lo, self.uc_start_hi,
             self.data_start_lo, self.data_start_hi) = struct.unpack_from('<9I', self.raw, 0x24)
        else:
            # gfx_firmware_header_v2_0: feature, ucode size/off, data size/off, start addr
            self.kind = 'gfx2'
            (_feat, self.mes_ucode_size, self.mes_ucode_offset,
             self.mes_data_size, self.mes_data_offset,
             self.uc_start_lo, self.uc_start_hi) = struct.unpack_from('<7I', self.raw, 0x20)
            self.mes_data_version = 0
            self.data_start_lo = self.data_start_hi = 0
        self.psp_hdr = None
        if self.raw[self.ucode_array_offset + 0x10:self.ucode_array_offset + 0x14] == b'$PS1':
            self.psp_hdr = self.raw[self.ucode_array_offset:self.ucode_array_offset + 0x100]
        self.uc_file = self.ucode_array_offset + (0x100 if self.psp_hdr else 0)
        self.uc = self.raw[self.uc_file:self.uc_file + self.mes_ucode_size]
        # code starts at uc_start_addr; the leading zero pad tells us the image base
        pad = 0
        while pad + 64 <= len(self.uc) and self.uc[pad:pad + 64] == bytes(64):
            pad += 64
        self.code_off = pad
        self.uc_base_va = (self.uc_start_lo - pad) & 0xffffffff
        self.data_file = self.ucode_array_offset + self.mes_data_offset
        self.data = self.raw[self.data_file:self.data_file + self.mes_data_size]
        # gfx v2 images give no data load address; their pointer tables live in the
        # ucode image itself, so point the data view at the ucode base instead.
        self.data_va = self.data_start_lo if self.kind == 'mes' else self.uc_base_va
        if self.kind != 'mes':
            self.data = self.uc
        self.md = cs.Cs(cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV64 | cs.CS_MODE_RISCVC)
        self.md.skipdata = True
        self._ins = None

    # ---- decode ---------------------------------------------------------
    def ins(self):
        if self._ins is None:
            self._ins = list(self.md.disasm(self.uc, self.uc_base_va))
        return self._ins

    def real(self):
        return [x for x in self.ins()
                if x.mnemonic not in ('.byte', '.db', 'unimp', 'c.unimp')]

    def code_lo(self):
        return self.uc_base_va + self.code_off

    def code_hi(self):
        return self.uc_base_va + len(self.uc)

    def word(self, va):
        """32-bit word at a VA in either image."""
        if self.data_va <= va < self.data_va + len(self.data) - 3:
            return struct.unpack_from('<I', self.data, va - self.data_va)[0]
        if self.uc_base_va <= va < self.uc_base_va + len(self.uc) - 3:
            return struct.unpack_from('<I', self.uc, va - self.uc_base_va)[0]
        return None

    def is_code(self, va):
        return self.code_lo() <= va < self.code_hi()

    # ---- analyses -------------------------------------------------------
    def functions(self):
        return [x.address for x in self.real()
                if x.mnemonic == 'addi' and x.op_str.startswith('sp, sp, -')]

    def pointer_tables(self, min_len=6):
        """Runs of consecutive dwords that are code addresses, in the data image."""
        runs, cur = [], []
        for i in range(0, len(self.data) - 3, 4):
            w = struct.unpack_from('<I', self.data, i)[0]
            if self.is_code(w):
                cur.append((self.data_va + i, w))
            else:
                if len(cur) >= min_len:
                    runs.append(cur)
                cur = []
        if len(cur) >= min_len:
            runs.append(cur)
        return runs

    def dispatches(self):
        """Every indirect jump/call, with the table base and bound if recoverable.

        Pattern this recognises (what the MES API dispatcher actually compiles to):
            lw    aX, off(pkt)          ; opcode out of the host packet
            li    aY, N
            bltu  aY, aX, default       ; opcode > N -> default
            lui   aT, hi ; addi aT, lo  ; table base
            slli  aX, aX, 2 ; add ; lw  ; jr
        """
        ins = self.ins()
        out = []
        for i, x in enumerate(ins):
            if x.mnemonic not in INDIRECT:
                continue
            lo = max(0, i - 24)
            win = ins[lo:i + 1]
            base = bound = None
            tbl_reg_hi = None
            for w in win:
                if w.mnemonic == 'lui':
                    m = re.match(r'(\w+), (0x[0-9a-f]+|\d+)', w.op_str)
                    if m:
                        tbl_reg_hi = (m.group(1), int(m.group(2), 0) << 12)
                elif w.mnemonic == 'addi' and tbl_reg_hi and w.op_str.startswith(tbl_reg_hi[0] + ','):
                    m = re.match(r'\w+, \w+, (-?0x[0-9a-f]+|-?\d+)', w.op_str)
                    if m:
                        base = tbl_reg_hi[1] + int(m.group(1), 0)
                elif w.mnemonic in ('li',):
                    m = re.match(r'(\w+), (-?0x[0-9a-f]+|-?\d+)', w.op_str)
                    if m:
                        bound = (m.group(1), int(m.group(2), 0))
                elif w.mnemonic in GUARDS:
                    if bound and bound[0] in w.op_str.split(',')[0]:
                        bound = ('%s %s' % (w.mnemonic, w.op_str), bound[1])
            if base is None and tbl_reg_hi:
                base = tbl_reg_hi[1]
            n = 0
            if base is not None:
                while n < 256:
                    v = self.word(base + 4 * n)
                    if v is None or not self.is_code(v):
                        break
                    n += 1
            guard = next((w for w in reversed(win) if w.mnemonic in GUARDS), None)
            out.append(dict(at=x.address, table=base, entries=n,
                            guard=('%s %s' % (guard.mnemonic, guard.op_str)) if guard else None,
                            bound=bound[1] if bound else None))
        return out


MEMCPY_VA = 0xf0021b28   # verified: xor/andi-7 alignment test, a7 = dst + len, word loop + byte tail
LOAD_MNEM = ('lw', 'lwu', 'lb', 'lbu', 'lh', 'lhu', 'ld', 'c.lw', 'c.ld', 'c.lwu')


def _ops(insn):
    return [o.strip() for o in insn.op_str.split(',')]


def provenance(img, ins, addrs, site, reg, back=200, depth=0):
    """Walk back linearly for the last writer of `reg` before `site`.

    Folds the RISC-V `lui`/`addi` constant pair, which is why site 0xf0014804 is a
    CONSTANT 0x2a154 and not a variable length -- an earlier version of this pass
    reported it as host-derived, which was wrong. Linear and intra-block only: it
    does not follow control flow, so treat 'unknown' as 'needs eyes', not as safe.
    """
    if depth > 4:
        return ('deep', '', None)
    i = addrs.index(site)
    for a in reversed(addrs[max(0, i - back):i]):
        x = ins[a]
        o = _ops(x)
        if not o or o[0] != reg:
            continue
        if x.mnemonic in ('li', 'c.li'):
            return ('const', int(o[1], 0), a)
        if x.mnemonic in ('addi', 'c.addi') and len(o) > 2:
            if o[1] == 'zero':
                return ('const', int(o[2], 0), a)
            base = provenance(img, ins, addrs, a, o[1], back, depth + 1)
            if base[0] == 'const':                       # lui/addi pair -> one constant
                return ('const', (base[1] + int(o[2], 0)) & 0xffffffffffffffff, a)
            return ('computed', x.op_str, a)
        if x.mnemonic == 'lui':
            return ('const', int(o[1], 0) << 12, a)
        if x.mnemonic in LOAD_MNEM:
            return ('load', x.op_str, a)
        if x.mnemonic in ('mv', 'c.mv'):
            return provenance(img, ins, addrs, a, o[1], back, depth + 1)
        return ('computed', '%s %s' % (x.mnemonic, x.op_str), a)
    return ('unknown', '', None)


def cmd_copies(img):
    """Classify every memcpy call site by the provenance of its length argument.

    This is the MES analogue of dmcub/dmcub_copyaudit.py, which is what found the
    DMCUB ABM_INIT_CONFIG overflow (FINDING-1): a u16 length loaded straight from a
    host frame into memcpy. Result for gc_11_0_0_mes_2: no such site exists.
    """
    ins = {x.address: x for x in img.ins()}
    addrs = sorted(ins)
    sites = []
    for a in addrs:
        x = ins[a]
        if x.mnemonic not in ('jal', 'call'):
            continue
        tgt = x.op_str.split()[-1]
        if not tgt.startswith('0x') or int(tgt, 16) != MEMCPY_VA:
            continue
        sites.append((a, provenance(img, ins, addrs, a, 'a2'),
                      provenance(img, ins, addrs, a, 'a0'),
                      provenance(img, ins, addrs, a, 'a1')))
    print('memcpy(0x%08x) call sites: %d' % (MEMCPY_VA, len(sites)))
    counts = {}
    for _, ln, _, _ in sites:
        counts[ln[0]] = counts.get(ln[0], 0) + 1
    print('length provenance: %s\n' % counts)
    for site, ln, dst, src in sites:
        if ln[0] == 'const':
            desc = 'len=0x%x const' % ln[1]
            flag = ''
        elif ln[0] == 'load':
            desc = 'len<-LOAD %s' % ln[1]
            flag = '   <== REVIEW: length from memory'
        else:
            desc = 'len<-%s %s' % (ln[0], str(ln[1])[:28])
            flag = '   <== REVIEW: length computed'
        print('  0x%08x  %-34s dst<-%-9s %s%s' % (
            site, desc, dst[0], str(dst[1])[:24], flag))
    var = [s for s in sites if s[1][0] != 'const']
    print('\n%d/%d sites use a constant length; %d need review' % (
        len(sites) - len(var), len(sites), len(var)))
    return sites


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('cmd', choices=['info', 'dispatch', 'tables', 'copies'])
    a = ap.parse_args()
    m = Mes(a.image)
    if a.cmd == 'info':
        print('%s' % a.image)
        print('  header %d.%d ip %d.%d, header_size 0x%x' % (
            m.hv_maj, m.hv_min, m.ip_maj, m.ip_min, m.hdr_size))
        print('  nested PSP header: %s' % ('$PS1 present' if m.psp_hdr else 'absent'))
        print('  ucode: file 0x%x size 0x%x -> image base VA 0x%08x, code at VA 0x%08x' % (
            m.uc_file, m.mes_ucode_size, m.uc_base_va, m.code_lo()))
        print('  data : file 0x%x size 0x%x -> VA 0x%08x' % (
            m.data_file, m.mes_data_size, m.data_va))
        ins, real = m.ins(), m.real()
        print('  decode: %d instructions, %d real (%.1f%%)' % (
            len(ins), len(real), 100.0 * len(real) / max(1, len(ins))))
        top = collections.Counter(x.mnemonic for x in real).most_common(8)
        print('  top mnemonics: %s' % ' '.join('%s:%d' % kv for kv in top))
        print('  function prologues (addi sp,sp,-N): %d' % len(m.functions()))
        print('  indirect jumps/calls: %d' % len(m.dispatches()))
    elif a.cmd == 'copies':
        cmd_copies(m)
    elif a.cmd == 'tables':
        for r in m.pointer_tables():
            print('  table VA 0x%08x  n=%-4d first: %s' % (
                r[0][0], len(r), ' '.join('0x%08x' % v for _, v in r[:4])))
    else:
        print('%-11s %-11s %-7s %-7s %s' % ('site', 'table', 'entries', 'bound', 'guard'))
        unguarded = 0
        for h in m.dispatches():
            if h['guard'] is None:
                unguarded += 1
            print('0x%08x %-11s %-7s %-7s %s' % (
                h['at'], ('0x%08x' % h['table']) if h['table'] else '-',
                h['entries'] or '-', h['bound'] if h['bound'] is not None else '-',
                h['guard'] or '*** none in window ***'))
        print('\n%d indirect sites, %d with no guard in the preceding window' % (
            len(m.dispatches()), unguarded))


if __name__ == '__main__':
    main()