"""Host-pointer taint double-fetch detector.

Closes the gap left by dfscan3: that pass only rooted pointers at a function's own
incoming arguments, so the most important pointer in the whole firmware -- the
MAPPED HOST BUFFER -- was invisible, because it arrives as an *out-parameter*:

    map_client_buffer(hi, lo, size, &handle, &out_va, ...);   // out_va <- host VA
    ...
    v = *out_va;            // reloaded from stack later -> opaque root in dfscan3

This tool seeds taint there instead:

  1. At a call to a known mapping function (`map_client_buffer`, `bl_map_rw_buffer`,
     `map_and_flush_range`, ...), every argument that is a *stack address* (SP+delta)
     is treated as a possible out-parameter, and that stack slot is marked HOST.
  2. Any register later loaded from a HOST-marked stack slot becomes HOST-rooted.
     `mov`/`add` propagate it, so `host+0x10` stays HOST at delta 0x10.
  3. Every load off a HOST root is recorded, with a flag for whether a compare
     follows shortly after (validated).
  4. Reported:
       (a) INTRA  - same HOST offset loaded twice with a compare between
                    => validate-then-reload on attacker memory
       (b) INTER  - a HOST pointer is passed to a callee, and the caller had
                    validated an offset the callee re-reads
                    => validate in caller, use in callee, on attacker memory

Anything reported here is, by construction, a double-fetch of memory the host owns.

Usage: python3 dfscan4.py <elf> <symbols.json>
"""
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]

MAPPERS = ('map_client_buffer','map_client_buffer_ro','map_and_flush_range',
           'map_page_and_flush','bl_map_rw_buffer','bl_map_ro_buffer',
           'shared_buffer_map','mbox_buffer_map','sram_window_map')


def scan(body, base_va, fva, fend, md, mapper_addrs, fname):
    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) < 3: return [], {}, []

    SP = arm.ARM_REG_SP
    ptr, nxt = {}, [0]
    host_slots = set()          # stack deltas holding a host pointer
    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)
        return ptr[r]

    host_acc = []               # (idx, off, mnem, addr)
    sp_addrs = set()            # stack deltas whose address was materialized
    validated = set()
    calls = []                  # (idx, tgt, {argidx:(root,delta)})
    hits = []

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

        if m in CMPS:
            for (j, off, _mn, _ad) in host_acc[-4:]:
                validated.add(off)
            continue

        if m.startswith('bl'):
            tgt=None
            try: tgt=int(x.op_str.replace('#','').strip(),16)
            except Exception: pass
            amap={ai: ptr[ar] for ai,ar in enumerate(ARGREGS) if ar in ptr}
            calls.append((i,tgt,amap))
            if tgt in mapper_addrs:
                # out-params may be REGISTER args (<=4) or STACK args (>4 arg fns
                # such as the 8-arg map_client_buffer). Over-approximate: any stack
                # slot whose address was taken is a candidate out-parameter.
                for ai,(root,delta) in amap.items():
                    if root=='SP': host_slots.add(delta)
                host_slots.update(sp_addrs)
            for ar in ARGREGS: ptr[ar]=fresh()
            continue

        # pc-relative literal -> PSP global
        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)
            if r=='SP': sp_addrs.add(d+ops[2].imm)   # &local passed as an out-param
            continue

        # load from a HOST-marked stack slot -> the loaded reg becomes HOST
        if m in LOADS_SIMPLE and len(ops)>=2 and ops[-1].type==arm.ARM_OP_MEM \
           and ops[0].type==arm.ARM_OP_REG:
            mem=ops[-1].mem
            if mem.base!=0 and mem.index==0:
                r,d=get(mem.base)
                if r=='SP' and (d+mem.disp) in host_slots:
                    ptr[ops[0].reg]=('HOST',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 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:
            if root=='HOST':
                # (a) INTRA: same host offset already read, and validated since
                for (j,poff,pmn,pad) in host_acc:
                    if poff==off and off in validated:
                        hits.append(('INTRA', fname, pad, x.address, off)); break
                host_acc.append((i,off,x.mnemonic,x.address))

        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()

    host_args = {}              # callee arg idx -> delta, for HOST pointers passed out
    for (i,tgt,amap) in calls:
        for ai,(root,delta) in amap.items():
            if root=='HOST':
                host_args.setdefault(tgt,[]).append((ai,delta,validated.copy()))
    return hits, host_args, calls


def main():
    elf,symf = sys.argv[1],sys.argv[2]
    d=json.load(open(symf)); raw=open(elf,'rb').read()
    ph=struct.unpack('<I',raw[28:32])[0]
    off,vaddr,paddr,fsz=struct.unpack('<IIII',raw[ph+4:ph+20]); body=raw[off:off+fsz]
    md=cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB); md.detail=True
    n2a={v.get('name'):int(k,16) for k,v in d['functions'].items()}
    mapper_addrs={n2a[n] for n in MAPPERS if n in n2a}
    print("mapping functions seeded: %s" % sorted(n for n in MAPPERS if n in n2a))

    # pass 1: per-function reads-off-arg summary (for the INTER check)
    argreads={}
    for k,v in d['functions'].items():
        fa=int(k,16); end=v.get('end')
        if not end: continue
        try: ins=list(md.disasm(body[fa-vaddr:end-vaddr], fa))
        except Exception: continue
        rd={}
        ptr={}
        for x in ins:
            m=x.mnemonic.split('.')[0]
            try: ops=x.operands
            except Exception: ops=[]
            if m in LOADS_SIMPLE and len(ops)>=2 and ops[-1].type==arm.ARM_OP_MEM:
                mem=ops[-1].mem
                if mem.base in ARGREGS and mem.index==0:
                    rd.setdefault(ARGREGS.index(mem.base),set()).add(mem.disp)
        if rd: argreads[fa]=rd

    allhits=[]; inter=[]
    names={int(k,16):v.get('name','?') for k,v in d['functions'].items()}
    for k,v in d['functions'].items():
        fa=int(k,16); end=v.get('end')
        if not end: continue
        h,hargs,_=scan(body,vaddr,fa,end,md,mapper_addrs,v.get('name','?'))
        allhits+=h
        for tgt,lst in hargs.items():
            for (ai,delta,val) in lst:
                for o in argreads.get(tgt,{}).get(ai,()):
                    if (o+delta) in val:
                        inter.append((v.get('name','?'), names.get(tgt,'?'), ai, o+delta))
    print("\nINTRA host double-fetches: %d"%len(allhits))
    for t,fn,a1,a2,off in allhits:
        print("   %-34s @%#08x then @%#08x  host_off=%#x"%(fn,a1,a2,off))
    print("\nINTER host double-fetches (validate in caller, re-read in callee): %d"%len(inter))
    seen=set()
    for cn,gn,ai,off in inter:
        if (cn,gn,off) in seen: continue
        seen.add((cn,gn,off))
        print("   %-30s validates host+%#x -> %-30s re-reads arg%d"%(cn,off,gn,ai))


if __name__=='__main__':
    main()