#!/usr/bin/env python3
"""Module-agnostic symbol porting + divergence detection across PSP firmware
builds.  Generalises vport.py: reads the load base from the ELF program header
instead of hardcoding sOS, and derives code ranges from the function table.

Naive same-offset comparison is WRONG (code shifts between builds and reports
hundreds of false changes), so matching is alignment-independent: index the old
image by N-mnemonic anchors, then verify the full normalised sequence.

  IDENT     found at same VA, bytes identical
  RELOC     found at same VA, mnemonics identical (immediates/addresses differ)
  MOVED     found at a different VA, logic unchanged
  DIVERGED  not found -> logic really changed, or function is new in the current build

usage: vport2.py <new.elf> <funcs.json> <old.elf> [...] [--list] [--only=SUBSTR]
"""
import json, sys, os, struct
from collections import defaultdict
from capstone import Cs, CS_ARCH_ARM, CS_MODE_THUMB, CS_MODE_LITTLE_ENDIAN

NGRAM = 8
md = Cs(CS_ARCH_ARM, CS_MODE_THUMB | CS_MODE_LITTLE_ENDIAN)
md.detail = False


def load(p):
    b = open(p, 'rb').read()
    phoff = struct.unpack('<I', b[28:32])[0]
    typ, off, va, pa, fsz, msz, fl, al = struct.unpack('<IIIIIIII', b[phoff:phoff + 32])
    return b, off, va, fsz


def seq(buf, va, maxn=None):
    out = []
    for ins in md.disasm(buf, va):
        out.append(ins.mnemonic)
        if maxn and len(out) >= maxn:
            break
    return tuple(out)


def seq_ins(buf, va):
    """(mnemonics, byte-offset-from-start) for each decoded instruction."""
    mn, offs = [], []
    for ins in md.disasm(buf, va):
        mn.append(ins.mnemonic); offs.append(ins.address - va)
    return mn, offs


def similarity(a, b):
    """fraction of positions whose mnemonics agree, over the longer sequence."""
    if not a or not b:
        return 0.0
    m = sum(1 for x, y in zip(a, b) if x == y)
    return m / float(max(len(a), len(b)))


FUZZ = 0.90   # >= this at best alignment => relocated, not changed


def ranges_from_funcs(F, gap=0x400):
    """Cluster function extents into contiguous code ranges."""
    rs = []
    cs, ce = None, None
    for st, nm, en in F:
        if cs is None:
            cs, ce = st, en
        elif st - ce > gap:
            rs.append((cs, ce)); cs, ce = st, en
        else:
            ce = max(ce, en)
    if cs is not None:
        rs.append((cs, ce))
    return rs


def build_index(ob, ooff, ova, ofsz, ranges):
    # NOTE: previously this indexed only windows derived from the CURRENT build's
    # function extents. Functions that moved BACKWARDS out of those windows were
    # then never found and reported DIVERGED (false positive) -- observed for the
    # heap2_* cluster, which sits at the end of SYS_DRV and shifted -0x1000.
    # Index the whole mapped image instead.
    idx = defaultdict(list)
    lo_lim, hi_lim = ova, ova + ofsz
    for lo, hi in [(lo_lim, hi_lim)]:
        for va in range(lo, hi, 2):
            fo = va - ova + ooff
            if fo + 64 > len(ob):
                continue
            s = seq(ob[fo:fo + 64], va, NGRAM)
            if len(s) == NGRAM:
                idx[s].append(va)
    return idx


def main():
    new_p, funcs_p = sys.argv[1], sys.argv[2]
    olds = [a for a in sys.argv[3:] if not a.startswith('--')]
    show = '--list' in sys.argv
    only = next((a.split('=', 1)[1] for a in sys.argv if a.startswith('--only=')), None)

    nb, noff, nva, nfsz = load(new_p)
    funcs = json.load(open(funcs_p))
    F = sorted(((int(k, 16), v[0], v[1]) for k, v in funcs.items()))
    ranges = ranges_from_funcs(F)

    for old_p in olds:
        ob, ooff, ova, ofsz = load(old_p)
        if ob == nb:
            continue
        idx = build_index(ob, ooff, ova, ofsz, ranges)
        counts = defaultdict(int)
        diverged = []
        for st, nm, en in F:
            if only and only not in nm:
                continue
            size = en - st
            fs = st - nva + noff
            if fs < 0 or fs + size > len(nb):
                continue
            body = nb[fs:fs + size]
            s_new = seq(body, st)
            if len(s_new) < NGRAM:
                counts['TINY'] += 1
                continue
            hit = None
            for cand in idx.get(s_new[:NGRAM], ()):
                co = cand - ova + ooff
                if co + size > len(ob):
                    continue
                if seq(ob[co:co + size], cand) == s_new:
                    hit = cand
                    break
            if hit is None:
                # FUZZY FALLBACK. Exact mnemonic equality gives false DIVERGEDs for
                # functions containing an inline literal pool: the pool holds a
                # relocated address, so its bytes differ between builds and are
                # disassembled as different "instructions" even though the code is
                # identical (observed: derive_psp_dram_hmac_key, heap2_*).
                # Try anchors at several offsets, then score the best alignment.
                mn_new, off_new = seq_ins(body, st)
                best = 0.0
                for k in range(0, max(1, min(len(mn_new) - NGRAM, 40)), 2):
                    anc = tuple(mn_new[k:k + NGRAM])
                    if len(anc) < NGRAM:
                        break
                    for cand in idx.get(anc, ()):
                        fs2 = cand - off_new[k]
                        co2 = fs2 - ova + ooff
                        if co2 < 0 or co2 + size > len(ob):
                            continue
                        b2 = seq(ob[co2:co2 + size], fs2)
                        sc = similarity(mn_new, list(b2))
                        if sc > best:
                            best = sc
                if best >= FUZZ:
                    counts['MOVED~'] += 1
                    continue
                counts['DIVERGED'] += 1
                diverged.append((nm, st, size, best))
            elif hit == st:
                counts['IDENT' if ob[st - ova + ooff:st - ova + ooff + size] == body else 'RELOC'] += 1
            else:
                counts['MOVED'] += 1
        print("=== %s   (%d compared)" % (os.path.basename(old_p), sum(counts.values())))
        print("    IDENT=%d RELOC=%d MOVED=%d TINY=%d  **DIVERGED=%d**" %
              (counts['IDENT'], counts['RELOC'], counts['MOVED'], counts['TINY'], counts['DIVERGED']))
        if show:
            for nm, st, sz, sc in sorted(diverged, key=lambda x: -x[2]):
                print("      DIVERGED  %-42s @0x%06x  %5d B  best_sim=%.2f" % (nm, st, sz, sc))
        print()


if __name__ == "__main__":
    main()