#!/usr/bin/env python3
"""vtable-anchored function map + disassembly report for Lofi-12XT (capstone).

Usage:
  python3 analysis/gen_funcmap.py "Lofi-12XT_v1.5.205/.../Lofi-12 XT.bin" /tmp/opencode/fw
Writes: funcmap.json, funcmap.md, asm-linear.txt into outdir.
"""
import os
import struct
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'tools'))
from lofi_image import parse_image, find_sect_by_role, code_file_off

from capstone import Cs, CS_ARCH_TMS320C64X, CS_MODE_LITTLE_ENDIAN


def collect_vfuncs(ro, code):
    p = ro['payload']
    words = [struct.unpack('<I', p[i:i + 4])[0]
             for i in range(0, len(p) - 3, 4)]
    is_ptr = [code['addr'] <= w < code['end'] for w in words]
    runs, i, n = [], 0, len(words)
    while i < n:
        if is_ptr[i]:
            j = i
            while j < n and is_ptr[j]:
                j += 1
            if j - i >= 3:
                runs.append((ro['addr'] + i * 4, words[i:j]))
            i = j
        else:
            i += 1
    funcs = {}
    for vaddr, ws in runs:
        for k, w in enumerate(ws):
            funcs.setdefault(w, []).append((vaddr, k))
    return runs, funcs


def main():
    image, outdir = sys.argv[1], sys.argv[2]
    os.makedirs(outdir, exist_ok=True)
    d = open(image, 'rb').read()
    info = parse_image(d)
    code = find_sect_by_role(info, 'code')
    ro = find_sect_by_role(info, 'rodata')
    entry = info['mmtd']['entry']
    runs, funcs = collect_vfuncs(ro, code)

    cs = Cs(CS_ARCH_TMS320C64X, CS_MODE_LITTLE_ENDIAN)
    rows = []


    def head_at(faddr):
        foff = code_file_off(info, faddr)
        if foff is None:
            return ["<outside code>"]
        for base_off, base_addr in ((foff, faddr),
                                    (foff - (faddr % 32), faddr - (faddr % 32))):
            if base_off < 0:
                continue
            try:
                insns = list(cs.disasm(d[base_off:base_off + 64], base_addr))
            except Exception as e:
                return ["<capstone err %s>" % e]
            for ins in insns:
                if ins.address == faddr or base_addr != faddr:
                    return ["%08X %s %s" % (i.address, i.mnemonic, i.op_str)
                            for i in insns[:6]] or ["<undecodable>"]
            if insns and base_addr == faddr:
                return ["%08X %s %s" % (i.address, i.mnemonic, i.op_str)
                        for i in insns[:6]]
        return ["<undecodable>"]

    for faddr in sorted(funcs):
        rows.append({"addr": faddr, "addr_hex": "%08X" % faddr,
                     "refs": funcs[faddr], "nrefs": len(funcs[faddr]),
                     "head": head_at(faddr)})
    # entry must be present even if not vtable-referenced
    if entry not in funcs:
        foff = code_file_off(info, entry)
        head = ["%08X %s %s" % (i.address, i.mnemonic, i.op_str)
                for i in list(cs.disasm(d[foff:foff + 64], entry))[:6]]
        rows.append({"addr": entry, "addr_hex": "%08X" % entry,
                     "refs": [], "nrefs": 0, "head": head})
        rows.sort(key=lambda r: r["addr"])

    import json
    json.dump({"image_size": len(d), "entry_hex": "%08X" % entry,
               "code": {"addr_hex": "%08X" % code['addr'], "len": code['len']},
               "rodata": {"addr_hex": "%08X" % ro['addr'], "len": ro['len']},
               "nvtable": len(runs), "nfunc": len(rows), "funcs": rows},
              open(os.path.join(outdir, "funcmap.json"), "w"), indent=1)

    with open(os.path.join(outdir, "funcmap.md"), "w") as f:
        f.write("# funcmap v1.5.205 — %d vtables, %d unique vfuncs (+entry %08X)\n\n"
                % (len(runs), len(rows), entry))
        f.write("| func addr | xrefs | first insn |\n|---|---|---|\n")
        for r in rows:
            f.write("| %s | %d | `%s` |\n"
                    % (r["addr_hex"], r["nrefs"], r["head"][0].replace("|", "/")))

    # linear sweep of whole code sect (fetch-packet granularity)
    with open(os.path.join(outdir, "asm-linear.txt"), "w") as f:
        cpay = code['payload']
        for i in range(0, len(cpay), 8):
            chunk = cpay[i:i + 8]
            if len(chunk) < 8:
                break
            for ins in cs.disasm(chunk, code['addr'] + i):
                f.write("%08X: %s %s\n" % (ins.address, ins.mnemonic, ins.op_str))
    print("vtables=%d funcs=%d entry=%08X" % (len(runs), len(rows), entry))
    print("wrote %s/{funcmap.json,funcmap.md,asm-linear.txt}" % outdir)


if __name__ == '__main__':
    main()