"""Extended ARM/Thumb double-fetch detector.

Improvements over dfscan.py (which only understood `ldr rX,[rB,#imm]` and keyed on
register *number*):

  1. LDRD  rX,rY,[rB,#off]      -> two accesses: off, off+4
  2. LDM   rB,{r1..rn}          -> n accesses: 0,4,...,4*(n-1)   (2754 were invisible)
  3. Indexed ldr [rB,rIdx]      -> access with unknown offset, tracked separately
                                   (792 were invisible; array/loop reads live here)
  4. POINTER IDENTITY instead of register number:
       mov  rD,rS        -> rD aliases rS's pointer
       add  rD,rS,#imm   -> rD = same root, delta+imm
     so `mov r5,r4 ; ldr [r5]` vs `ldr [r4]` is now recognised as the same location,
     and `add r7,r4,#8 ; ldm r7` is attributed to root r4 at offsets 8,12,...

A hit = the same (pointer-root, offset) is loaded 2+ times, with a compare/branch
between the loads, and the root pointer is not redefined in between.

Usage: python3 dfscan2.py <elf> <symbols.json> [--indexed]
"""
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','it','ite','itt','itte','ittt'}


def scan_func(body, base_va, fva, fend, name, md, want_indexed):
    start, n = fva - base_va, fend - fva
    if start < 0 or n <= 0 or start + n > len(body):
        return []
    try:
        ins = list(md.disasm(body[start:start+n], fva))
    except Exception:
        return []
    if len(ins) < 4:
        return []

    # pointer identity: reg -> (root_id, delta). fresh root when written opaquely.
    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 ('r%d' % r), 0) if r == SP else (('r%d' % r), 0)
            if r == SP: ptr[r] = ('SP', 0)
        return ptr[r]

    accesses = {}   # (root, off) -> [indices]
    cmp_idx = []
    root_redef = {} # root -> [indices where redefined]

    for i, x in enumerate(ins):
        m = x.mnemonic.split('.')[0]
        if m in CMPS:
            cmp_idx.append(i)

        # --- track pointer moves BEFORE recording the access ---
        try:
            ops = x.operands
        except Exception:
            ops = []

        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

        # --- record loads ---
        # a register loaded from a pc-relative literal holds a PSP GLOBAL address
        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

        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:
                r, d = get(mem.base)
                if mem.index != 0:
                    if want_indexed: recs.append((r, 'IDX'))
                else:
                    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:
            if ops[0].type == arm.ARM_OP_REG:
                r, d = get(ops[0].reg)
                nregs = max(0, len(ops) - 1)
                recs += [(r, d + 4*k) for k in range(nregs)]

        for key in recs:
            accesses.setdefault(key, []).append(i)

        # --- writes invalidate pointer identity (after the access is recorded) ---
        try:
            _, wr = x.regs_access()
        except Exception:
            wr = []
        for w in wr:
            if m == 'mov' or m in ('add','addw'):
                continue
            if w == SP:
                ptr[w] = ('SP', 0)      # stack stays the stack: sub sp,#N must not re-root
                continue
            ptr[w] = fresh()
            root_redef.setdefault(('r%d' % w), []).append(i)

    hits = []
    for (root, off), idxs in accesses.items():
        if len(idxs) < 2:
            continue
        for a, b in zip(idxs, idxs[1:]):
            if not any(a < c < b for c in cmp_idx):
                continue
            if any(a < c < b for c in root_redef.get(root, [])):
                continue
            if root in ('SP','GLOBAL') and '--stack' not in sys.argv:
                break   # stack spills and PSP globals (e.g. the stack canary) are not host memory
            hits.append((name, fva, ins[a].address, ins[b].address, root, off, b - a,
                         ins[a].mnemonic, ins[b].mnemonic))
            break
    return hits


def main():
    elf, symf = sys.argv[1], sys.argv[2]
    want_indexed = '--indexed' in sys.argv
    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_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

    all_hits = []
    for k, v in d['functions'].items():
        fva = int(k, 16); end = v.get('end')
        if not end: continue
        best = []
        for md in (md_t, md_a):
            h = scan_func(body, vaddr, fva, end, v.get('name','?'), md, want_indexed)
            if len(h) > len(best): best = h
        all_hits += best
    print("functions scanned: %d   double-fetch candidates: %d"
          % (len(d['functions']), len(all_hits)))
    for nm, fva, a1, a2, root, offv, gap, m1, m2 in sorted(all_hits, key=lambda x: x[6]):
        o = offv if offv == 'IDX' else ('%#x' % offv)
        print("  %-38s func=%#08x  @%#08x(%s),%#08x(%s)  base=%s off=%s gap=%d"
              % (nm, fva, a1, m1, a2, m2, root, o, gap))


if __name__ == '__main__':
    main()