#!/usr/bin/env python3
"""Identify the DMCUB core ISA empirically: disassemble the code region under
every plausible capstone architecture and score linear-sweep coverage.

Scoring: sweep N bytes from a fixed offset, count bytes consumed by successfully
decoded instructions vs bytes skipped, and additionally require the decode to be
"self-consistent" (a wrong ISA typically decodes a low fraction and produces a
scatter of rare mnemonics).
"""
import collections
import sys

import capstone as cs

ARCHS = [
    ('riscv32', cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV32),
    ('riscv32c', cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV32 | cs.CS_MODE_RISCVC),
    ('riscv64', cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV64),
    ('riscv64c', cs.CS_ARCH_RISCV, cs.CS_MODE_RISCV64 | cs.CS_MODE_RISCVC),
    ('arm', cs.CS_ARCH_ARM, cs.CS_MODE_ARM),
    ('thumb', cs.CS_ARCH_ARM, cs.CS_MODE_THUMB),
    ('arm64', cs.CS_ARCH_ARM64, cs.CS_MODE_ARM),
    ('mips32el', cs.CS_ARCH_MIPS, cs.CS_MODE_MIPS32 | cs.CS_MODE_LITTLE_ENDIAN),
    ('ppc32', cs.CS_ARCH_PPC, cs.CS_MODE_32),
    ('sparc', cs.CS_ARCH_SPARC, cs.CS_MODE_BIG_ENDIAN),
    ('x86-32', cs.CS_ARCH_X86, cs.CS_MODE_32),
]


def score(data, arch, mode, base=0):
    try:
        md = cs.Cs(arch, mode)
    except cs.CsError:
        return None
    md.skipdata = True
    covered = 0
    skipped = 0
    mnem = collections.Counter()
    for ins in md.disasm(data, base):
        if ins.mnemonic in ('.byte', '.db'):
            skipped += ins.size
        else:
            covered += ins.size
            mnem[ins.mnemonic] += 1
    total = covered + skipped
    return {
        'cover': covered / total if total else 0.0,
        'ninsn': sum(mnem.values()),
        'uniq': len(mnem),
        'top': mnem.most_common(8),
    }


def main():
    path = sys.argv[1]
    off = int(sys.argv[2], 0) if len(sys.argv) > 2 else 0x100
    n = int(sys.argv[3], 0) if len(sys.argv) > 3 else 0x4000
    data = open(path, 'rb').read()[off:off + n]
    print('%s @ 0x%x len 0x%x' % (path, off, len(data)))
    rows = []
    for name, arch, mode in ARCHS:
        r = score(data, arch, mode)
        if r:
            rows.append((r['cover'], name, r))
    for cover, name, r in sorted(rows, reverse=True):
        print('  %-9s cover=%5.1f%%  insns=%-6d uniq=%-4d %s' % (
            name, cover * 100, r['ninsn'], r['uniq'],
            ' '.join('%s:%d' % (m, c) for m, c in r['top'])))


if __name__ == '__main__':
    main()