"""Interprocedural double-fetch detector.

dfscan/dfscan2 are intra-procedural: they only see a validate-then-reload that happens
inside ONE function. The remaining structural class is the split:

    F:  x = P->len;  if (x > LIMIT) reject;      // validate, in F
        G(P);                                     // pass the SAME pointer down
    G:  memcpy(dst, P->buf, P->len);              // use, re-read in G

Neither function alone looks wrong. This tool finds them.

How it works
------------
1. SUMMARISE every function: with pointer-identity tracking (mov/add propagation),
   record for each *incoming argument* r0-r3 the set of offsets read off it, and
   which of those reads are followed by a compare ("validated").
       summary[f] = { argi: {offsets read} },  validated[f] = { argi: {offsets} }
2. Propagate summaries transitively (fixpoint): if F calls G passing its own arg,
   G's reads become F's reads at the translated offset.
3. REPORT: at a call site F -> G where F passes pointer P (rooted at one of F's own
   args, at delta d) into G's argument i, and
       F validated (P, d+o)  BEFORE the call, and
       G reads      (arg i, o)
   then the value validated in F is re-read in G => interprocedural double-fetch.

Only argument-rooted pointers are considered: stack ('SP') and pc-relative globals
('GLOBAL') are PSP-internal and cannot be mutated by the host mid-call.

Usage: python3 dfscan3.py <elf> <symbols.json> [--maxdepth N]
"""
import json, struct, sys
import capstone as cs
from capstone import arm

LOADS_SIMPLE = {'ldr','ldrb','ldrh','ldrsb','ldrsh'}
LOAD_PAIR    = {'ldrd'}
LOAD_MULTI   = {'ldm','ldmia','ldmdb','ldmib','ldmda'}
CMPS         = {'cmp','cmn','tst','teq','subs','cbz','cbnz'}
ARGREGS      = [arm.ARM_REG_R0, arm.ARM_REG_R1, arm.ARM_REG_R2, arm.ARM_REG_R3]


def analyse(body, base_va, fva, fend, md):
    """Return (reads, validated, calls) for one function.
    reads/validated: {argidx: set(offsets)}
    calls: list of (insn_index, target, {argidx: (argroot, delta)})"""
    start, n = fva - base_va, fend - fva
    if start < 0 or n <= 0 or start + n > len(body):
        return None
    try:
        ins = list(md.disasm(body[start:start+n], fva))
    except Exception:
        return None
    if len(ins) < 3:
        return None

    SP = arm.ARM_REG_SP
    ptr, nxt = {}, [0]
    def fresh():
        nxt[0] += 1
        return ('t%d' % nxt[0], 0)
    def get(r):
        if r not in ptr:
            ptr[r] = ('SP', 0) if r == SP else ('a%d' % ARGREGS.index(r), 0) if r in ARGREGS else ('r%d' % r, 0)
        return ptr[r]

    reads, validated, calls = {}, {}, []
    acc_seq = []           # (index, root, off)

    for i, x in enumerate(ins):
        m = x.mnemonic.split('.')[0]
        try: ops = x.operands
        except Exception: ops = []

        if m in CMPS:
            # mark the most recent accesses as validated
            for (j, root, off) in acc_seq[-4:]:
                if root.startswith('a'):
                    validated.setdefault(int(root[1:]), set()).add(off)
            continue

        if m.startswith('bl'):
            tgt = None
            try: tgt = int(x.op_str.replace('#','').strip(), 16)
            except Exception: pass
            amap = {}
            for ai, ar in enumerate(ARGREGS):
                if ar in ptr:
                    amap[ai] = ptr[ar]
            calls.append((i, tgt, amap))
            for ar in ARGREGS:                      # call clobbers r0-r3
                ptr[ar] = fresh()
            continue

        if m in LOADS_SIMPLE and len(ops) >= 2 and ops[-1].type == arm.ARM_OP_MEM \
           and ops[-1].mem.base == arm.ARM_REG_PC and ops[0].type == arm.ARM_OP_REG:
            ptr[ops[0].reg] = ('GLOBAL', 0); continue
        if m == 'mov' and len(ops) == 2 and ops[0].type == arm.ARM_OP_REG and ops[1].type == arm.ARM_OP_REG:
            ptr[ops[0].reg] = get(ops[1].reg); continue
        if m in ('add','addw') and len(ops) == 3 and ops[0].type == arm.ARM_OP_REG \
           and ops[1].type == arm.ARM_OP_REG and ops[2].type == arm.ARM_OP_IMM:
            r, d = get(ops[1].reg); ptr[ops[0].reg] = (r, d + ops[2].imm); continue

        recs = []
        if m in LOADS_SIMPLE and len(ops) >= 2 and ops[-1].type == arm.ARM_OP_MEM:
            mem = ops[-1].mem
            if mem.base != 0 and mem.index == 0:
                r, d = get(mem.base); recs.append((r, d + mem.disp))
        elif m in LOAD_PAIR and len(ops) >= 3 and ops[-1].type == arm.ARM_OP_MEM:
            mem = ops[-1].mem
            if mem.base != 0 and mem.index == 0:
                r, d = get(mem.base); recs += [(r, d+mem.disp), (r, d+mem.disp+4)]
        elif m in LOAD_MULTI and ops and ops[0].type == arm.ARM_OP_REG:
            r, d = get(ops[0].reg)
            recs += [(r, d + 4*k) for k in range(max(0, len(ops)-1))]

        for (root, off) in recs:
            acc_seq.append((i, root, off))
            if root.startswith('a'):
                reads.setdefault(int(root[1:]), set()).add(off)

        try: _, wr = x.regs_access()
        except Exception: wr = []
        for w in wr:
            if m in ('mov','add','addw'): continue
            ptr[w] = ('SP', 0) if w == SP else fresh()
    return reads, validated, calls


def main():
    elf, symf = sys.argv[1], sys.argv[2]
    d = json.load(open(symf))
    raw = open(elf, 'rb').read()
    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]
    md = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB); md.detail = True

    info, names = {}, {}
    for k, v in d['functions'].items():
        fva = int(k, 16); end = v.get('end')
        names[fva] = v.get('name', '?')
        if not end: continue
        r = analyse(body, vaddr, fva, end, md)
        if r: info[fva] = r

    # transitive read propagation to fixpoint
    reads = {f: {a: set(o) for a, o in info[f][0].items()} for f in info}
    for _ in range(4):
        changed = False
        for f, (_, _, calls) in info.items():
            for (_, tgt, amap) in calls:
                if tgt not in reads: continue
                for ai, (root, delta) in amap.items():
                    if not root.startswith('a'): continue
                    caller_arg = int(root[1:])
                    for off in reads[tgt].get(ai, ()):
                        newoff = off + delta
                        if newoff not in reads.setdefault(f, {}).setdefault(caller_arg, set()):
                            reads[f][caller_arg].add(newoff); changed = True
        if not changed: break

    hits = []
    for f, (_, validated, calls) in info.items():
        for (idx, tgt, amap) in calls:
            if tgt not in reads: continue
            for ai, (root, delta) in amap.items():
                if not root.startswith('a'): continue
                caller_arg = int(root[1:])
                vset = validated.get(caller_arg, set())
                if not vset: continue
                for off in reads[tgt].get(ai, ()):
                    if off + delta in vset:
                        hits.append((names.get(f,'?'), f, names.get(tgt,'?'), tgt,
                                     caller_arg, ai, off + delta, delta))
    print("functions analysed: %d   interprocedural double-fetch candidates: %d"
          % (len(info), len(hits)))
    seen = set()
    for cn, cf, gn, gf, ca, ga, off, delta in sorted(hits, key=lambda x: x[0]):
        key = (cn, gn, off)
        if key in seen: continue
        seen.add(key)
        print("  %-30s validates arg%d+%#x  ->  %-30s re-reads arg%d+%#x"
              % (cn, ca, off, gn, ga, off - delta))


if __name__ == '__main__':
    main()