#!/usr/bin/env python3
"""Locate a DIVERGED function in an older build and print a side-by-side
instruction diff.

A diverged function cannot be found by mnemonic match (that is what made it
diverged), so we locate it positionally: find the nearest surrounding functions
that DID match, derive the local code shift from them, and disassemble the old
image at the implied address.

usage: vdiff_func.py <new.elf> <funcs.json> <old.elf> <func_name>
"""
import json, sys, os
from collections import defaultdict
from capstone import Cs, CS_ARCH_ARM, CS_MODE_THUMB, CS_MODE_LITTLE_ENDIAN

OFF = 0x1000
NGRAM = 8
CODE_RANGES = [(0x0, 0x7022), (0x12000, 0x17dd0)]

md = Cs(CS_ARCH_ARM, CS_MODE_THUMB | CS_MODE_LITTLE_ENDIAN)


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):
    idx = defaultdict(list)
    for lo, hi in CODE_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, old_p, target = sys.argv[1:5]
    new = open(new_p, 'rb').read()
    old = open(old_p, 'rb').read()
    funcs = json.load(open(funcs_p))
    F = sorted(((int(k, 16), v[0], v[1]) for k, v in funcs.items()))
    idx = build_index(old)

    # map every matchable function to its old address -> local shift
    shifts = {}
    tgt = None
    for st, nm, en in F:
        if nm == target:
            tgt = (st, en)
        size = en - st
        body = new[st + OFF:st + OFF + size]
        s_new = seq(body, st)
        if len(s_new) < NGRAM:
            continue
        for cand in idx.get(s_new[:NGRAM], ()):
            if seq(old[cand + OFF:cand + OFF + size], cand) == s_new:
                shifts[st] = cand - st
                break
    if tgt is None:
        print("function not found:", target); return
    st, en = tgt
    before = [a for a in shifts if a < st]
    after = [a for a in shifts if a > st]
    db = shifts[max(before)] if before else 0
    da = shifts[min(after)] if after else db
    print("target %s @0x%05x size=%d" % (target, st, en - st))
    print("nearest matched before: 0x%05x shift=%+d" % (max(before), db) if before else "none")
    print("nearest matched after : 0x%05x shift=%+d" % (min(after), da) if after else "none")

    size = en - st
    for shift in sorted({db, da}):
        oa = st + shift
        print("\n===== OLD @0x%05x (shift %+d) vs NEW @0x%05x =====" % (oa, shift, st))
        A = list(md.disasm(new[st + OFF:st + OFF + size], st))
        B = list(md.disasm(old[oa + OFF:oa + OFF + size + 16], oa))
        for i in range(max(len(A), len(B))):
            l = "%-9s %-7s %s" % (hex(A[i].address), A[i].mnemonic, A[i].op_str) if i < len(A) else ""
            r = "%-9s %-7s %s" % (hex(B[i].address), B[i].mnemonic, B[i].op_str) if i < len(B) else ""
            mark = ""
            if i < len(A) and i < len(B) and A[i].mnemonic != B[i].mnemonic:
                mark = "  <<<"
            print("%-46s | %-46s%s" % (l, r, mark))


if __name__ == "__main__":
    main()