#!/usr/bin/env python3
"""Port current-build function symbols onto older PSP firmware builds and
report where they DIVERGE.

Naive same-offset comparison is WRONG: code shifts between builds, so a function
at VA X in the new build sits at VA X-k in the old one, and a byte/mnemonic
compare at the same offset comes back "changed" for hundreds of functions that
are actually identical.  (Observed: sos .51-vs-.54 reported 256 "LOGIC" changes;
spot-checking showed the sequences were identical, merely shifted 0x40 bytes.)

So we do alignment-independent matching:
  1. index the OLD image: for every 2-byte-aligned offset in its code ranges,
     hash the next N mnemonics;
  2. for each named function in the NEW build, look up its first-N mnemonics to
     get candidate offsets, then verify the FULL normalised mnemonic sequence.

Classification per function:
  IDENT   found in old build, byte-identical
  RELOC   found, mnemonics identical, bytes differ (immediates/addresses only)
  MOVED   found at a different address (still counts as unchanged logic)
  DIVERGED  not found anywhere -> logic really changed, or function is new

Only DIVERGED is interesting for vulnerability archaeology.

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

OFF = 0x1000          # mkelf.py DATA_OFF: file offset of VA 0
NGRAM = 8             # anchor length in instructions
CODE_RANGES = [(0x0, 0x7022), (0x12000, 0x17dd0)]   # sOS CODE1, CODE2 (VA)

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


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 build_index(old, ranges):
    idx = defaultdict(list)
    for lo, hi in ranges:
        for va in range(lo, hi, 2):
            fo = va + OFF
            if fo + 64 > len(old):
                continue
            s = seq(old[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
    new = open(new_p, 'rb').read()
    funcs = json.load(open(funcs_p))
    F = sorted(((int(k, 16), v[0], v[1]) for k, v in funcs.items()))

    for old_p in olds:
        old = open(old_p, 'rb').read()
        if old == new:
            continue
        idx = build_index(old, CODE_RANGES)
        counts = defaultdict(int)
        diverged = []
        for st, nm, en in F:
            size = en - st
            fs = st + OFF
            if fs + size > len(new):
                continue
            body = new[fs:fs + size]
            s_new = seq(body, st)
            if len(s_new) < NGRAM:
                counts['TINY'] += 1
                continue
            anchor = s_new[:NGRAM]
            hit = None
            for cand in idx.get(anchor, ()):
                co = cand + OFF
                if co + size > len(old):
                    continue
                if seq(old[co:co + size], cand) == s_new:
                    hit = cand
                    break
            if hit is None:
                counts['DIVERGED'] += 1
                diverged.append((nm, st, size))
            elif hit == st:
                counts['IDENT' if old[fs:fs + size] == body else 'RELOC'] += 1
            else:
                counts['MOVED'] += 1
        ver = os.path.basename(old_p)
        tot = sum(counts.values())
        print("=== %s   (%d functions compared)" % (ver, tot))
        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 in sorted(diverged, key=lambda x: -x[2]):
                print("      DIVERGED  %-42s @0x%05x  %d B" % (nm, st, sz))
        print()


if __name__ == "__main__":
    main()