#!/usr/bin/env python3
"""Function inventory + indirect-jump (switch table) discovery for a DMCUB image.

Requires capstone >= 6 (Xtensa support). On this box:
    /home/user/.venvs/cs6/bin/python

Layout facts this relies on (established empirically, see AUDIT-DMCUB-FIRMWARE.md):
  * payload starts at file offset 0x100 and is loaded at VA 0x60000000 (DMUB CW0)
  * literal pools sit in the low part of the image, code is contiguous after it,
    so sweeping ENTRY_i .. ENTRY_{i+1} recovers function bodies without needing
    a full CFG.

Emits:
  * function count / sizes
  * every `jx` (indirect jump = switch dispatch) with the l32r-loaded table
    address in scope, and the decoded table if it lands in the image
  * per-function references to the DMCUB cache windows (0x60..0x66 000000)
"""
import argparse
import collections
import struct
import sys

import capstone as cs

BASE = 0x60000000
PAYLOAD_OFF = 0x100


class Image:
    """A DMCUB container, with the loadable body isolated and mapped.

    Layout (verified, see AUDIT-DMCUB-02.md):
        0x000  dmcub_firmware_header_v1_0
        0x100  **PSP header** ($PS1 magic at +0x10), 0x100 bytes
        0x200  body -- this is what the driver copies to DMUB_CW0_BASE
               (amdgpu: `fw_inst_const + PSP_HEADER_BYTES`, size also reduced
                by PSP_HEADER_BYTES)
        end-256  PSP footer / signature

    So VA = 0x60000000 + (file_off - 0x200). An earlier pass had this 0x100 low;
    the decisive check is that the firmware's own register-table base literal
    0x60018030 then lands exactly on the register-address table.
    """

    def __init__(self, path, payload_off=None, base=BASE):
        self.raw = open(path, 'rb').read()
        self.path = path
        self.base = base
        (size_bytes, hdr_size, _, _, _, _, _, ucode_size, ucode_off,
         _) = struct.unpack_from('<IIHHHHIIII', self.raw, 0)
        body = ucode_off
        self.psp_header = None
        if self.raw[ucode_off + 0x10:ucode_off + 0x14] in (b'$PS1', b'$PS2'):
            self.psp_header = self.raw[ucode_off:ucode_off + 0x100]
            body = ucode_off + 0x100
            ucode_size -= 0x100
        self.off = body if payload_off is None else payload_off
        self.payload = self.raw[self.off:self.off + ucode_size]
        self.md = cs.Cs(cs.CS_ARCH_XTENSA, cs.CS_MODE_XTENSA_ESP32)
        self.md.detail = True
        # Without skipdata a single undecodable byte silently truncates the rest
        # of the function -- which cost one wrong "only 539 literals" result.
        self.md.skipdata = True

    def va(self, poff):
        return self.base + poff

    def poff(self, va):
        return va - self.base

    def inside(self, va):
        return 0 <= va - self.base < len(self.payload)

    def word(self, va):
        p = self.poff(va)
        if 0 <= p < len(self.payload) - 3:
            return struct.unpack_from('<I', self.payload, p)[0]
        return None

    def code_start(self, ents=None, max_gap=0x1000):
        """Start of the code region: the beginning of the longest run of `entry`
        prologues spaced less than `max_gap` apart.

        The low part of an image is literal pools and tables, which yield stray
        `36 41 xx` matches -- locally dense enough to fool a simple
        density-window rule, which put the code start at 0x4e8 and inflated the
        Renoir/DCN function counts. The longest-run rule puts it in the right
        place for every image in fw/."""
        ents = self.entries() if ents is None else ents
        if not ents:
            return 0
        best = (0, ents[0])
        run_start, run_n = ents[0], 1
        for prev, cur in zip(ents, ents[1:]):
            if cur - prev < max_gap:
                run_n += 1
            else:
                if run_n > best[0]:
                    best = (run_n, run_start)
                run_start, run_n = cur, 1
        if run_n > best[0]:
            best = (run_n, run_start)
        return best[1]

    def entries(self):
        """Offsets of `entry a1, N` prologues, validated through capstone."""
        out = []
        p = self.payload
        for i in range(len(p) - 3):
            if p[i] == 0x36 and (p[i + 1] & 0x0F) == 0x01:
                ins = next(self.md.disasm(p[i:i + 3], self.va(i), 1), None)
                if ins and ins.mnemonic == 'entry':
                    out.append(i)
        return out


def analyse(img):
    ents = img.entries()
    lo = img.code_start(ents)
    ents = [e for e in ents if e >= lo]
    funcs = []
    jx_sites = []
    win_refs = collections.defaultdict(collections.Counter)
    for idx, start in enumerate(ents):
        end = ents[idx + 1] if idx + 1 < len(ents) else len(img.payload)
        code = img.payload[start:end]
        fva = img.va(start)
        ninsn = 0
        lit_scope = collections.deque(maxlen=8)   # recent l32r values, for jx tables
        for ins in img.md.disasm(code, fva):
            ninsn += 1
            m = ins.mnemonic
            if m == 'l32r':
                # operand is the literal *address*; resolve one level
                lit_addr = None
                for op in ins.operands:
                    if op.type == cs.xtensa.XTENSA_OP_MEM:
                        lit_addr = op.mem.disp
                    elif op.type == cs.xtensa.XTENSA_OP_IMM:
                        lit_addr = op.imm
                if lit_addr is None:
                    txt = ins.op_str.split()[-1]
                    try:
                        lit_addr = int(txt, 0)
                    except ValueError:
                        lit_addr = None
                if lit_addr is not None:
                    v = img.word(lit_addr)
                    if v is not None:
                        lit_scope.append((ins.address, v))
                        hb = v >> 24
                        if 0x60 <= hb <= 0x66:
                            win_refs[hb][fva] += 1
            elif m == 'jx':
                jx_sites.append((ins.address, fva, list(lit_scope)))
        funcs.append((fva, end - start, ninsn))
    return ents, funcs, jx_sites, win_refs


def decode_table(img, addr, limit=256):
    """Read a candidate switch table: consecutive words that are code VAs."""
    out = []
    a = addr
    while len(out) < limit:
        v = img.word(a)
        if v is None:
            break
        if not img.inside(v) or img.poff(v) < 0x1c000:
            break
        out.append(v)
        a += 4
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('--tables', action='store_true')
    ap.add_argument('--windows', action='store_true')
    args = ap.parse_args()

    img = Image(args.image)
    ents, funcs, jx_sites, win_refs = analyse(img)
    sizes = [f[1] for f in funcs]
    print('%s  payload 0x%x bytes @ VA 0x%08x' % (args.image, len(img.payload), img.base))
    print('windowed functions (entry a1): %d' % len(ents))
    print('code span: VA 0x%08x .. 0x%08x' % (img.va(ents[0]), img.va(ents[-1])))
    print('function size: median %d  max %d  total %d bytes' % (
        sorted(sizes)[len(sizes) // 2], max(sizes), sum(sizes)))
    print('indirect jumps (jx): %d' % len(jx_sites))

    if args.windows:
        print('\ncache-window references by literal high byte:')
        for hb in sorted(win_refs):
            fns = win_refs[hb]
            print('  0x%02x000000: %d refs in %d functions; top: %s' % (
                hb, sum(fns.values()), len(fns),
                ' '.join('%08x:%d' % (f, n) for f, n in fns.most_common(5))))

    if args.tables:
        print('\njx sites and candidate tables:')
        for site, fva, lits in jx_sites:
            tabs = []
            for la, v in lits:
                t = decode_table(img, v)
                if len(t) >= 4:
                    tabs.append((v, t))
            desc = ', '.join('table@0x%08x n>=%d' % (v, len(t)) for v, t in tabs) or 'no table in scope'
            print('  jx @0x%08x (func 0x%08x): %s' % (site, fva, desc))


if __name__ == '__main__':
    main()