"""ASM-level double-fetch scanner.

Signature: the SAME [base, #off] is loaded twice within a function, with a
comparison/branch between the two loads -- i.e. validate-then-reload. If `base`
holds a pointer to attacker memory, the second load can see different data.
"""
import json, struct, sys, capstone as cs
from collections import defaultdict

elf = sys.argv[1]; symf = sys.argv[2]
d = json.load(open(symf))
raw = open(elf,'rb').read()
# minimal ELF32 PT_LOAD
phoff = struct.unpack('<I', raw[28:32])[0]
off, vaddr, paddr, fsz = struct.unpack('<IIII', raw[phoff+4:phoff+20])
body = raw[off:off+fsz]
base_va = vaddr

md_t = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB); md_t.detail = True
md_a = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_ARM);   md_a.detail = True

LOADS = {'ldr','ldrb','ldrh','ldrsb','ldrsh'}
CMPS  = {'cmp','cmn','tst','teq','subs','cbz','cbnz'}

def scan(fva, fend, name):
    start = fva - base_va
    n = fend - fva
    if start < 0 or n <= 0 or start+n > len(body): return None
    code = body[start:start+n]
    best = None
    for md in (md_t, md_a):
        seq = []
        try: ins_list = list(md.disasm(code, fva))
        except Exception: continue
        if len(ins_list) < 6: continue
        loads = defaultdict(list)   # (basereg, off) -> [index]
        cmp_idx = []
        for i, ins in enumerate(ins_list):
            m = ins.mnemonic.split('.')[0]
            if m in CMPS: cmp_idx.append(i)
            if m in LOADS and len(ins.operands) >= 2:
                o = ins.operands[1]
                if o.type == cs.arm.ARM_OP_MEM and o.mem.base != 0 and o.mem.index == 0:
                    loads[(o.mem.base, o.mem.disp)].append(i)
        for (breg, disp), idxs in loads.items():
            if len(idxs) < 2: continue
            for a, b in zip(idxs, idxs[1:]):
                if any(a < c < b for c in cmp_idx):
                    # base reg must not be written between the loads (else diff ptr)
                    rewritten = False
                    for k in range(a+1, b):
                        for w in ins_list[k].regs_access()[1]:
                            if w == breg: rewritten = True; break
                        if rewritten: break
                    if not rewritten:
                        rec = (name, fva, ins_list[a].address, ins_list[b].address,
                               ins_list[a].reg_name(breg), disp, b-a)
                        if best is None or (b-a) < best[6]: best = rec
    return best

funcs = d['functions']
hits = []
for k, v in funcs.items():
    fva = int(k, 16) if not k.isdigit() else int(k)
    try: fva = int(k, 16)
    except ValueError: fva = int(k)
    end = v.get('end')
    if not end: continue
    r = scan(fva, end, v.get('name','?'))
    if r: hits.append(r)
print("functions scanned: %d   double-fetch candidates: %d" % (len(funcs), len(hits)))
for nm, fva, a1, a2, breg, disp, gap in sorted(hits, key=lambda x: x[6]):
    print("  %-34s func=%#07x  loads @%#07x,%#07x  [%s,#%d]  gap=%d" % (nm, fva, a1, a2, breg, disp, gap))