#!/usr/bin/env python3
"""Which ARM-mode entry points are reachable by clearing the low byte of a LIVE saved LR?

The LSB-zero trick (FE09 zero-store aimed at P = LR_slot-3, Add.6/7) can only redirect a return
whose frame is live AT THE MOMENT OF THE WRITE, because any later PUSH rewrites the whole word.
So the reachable ARM-mode entry set is

    { retaddr & ~0xFF : retaddr is a call site in a function on the chain containing the write }

and the question is whether that set intersects the ARM-mode pivots from arm_mode_sled_sweep.py.

METHOD (deliberately an OVER-approximation, so an empty intersection is a sound negative):
  * global sliding-window scan at every even address for BL/BLX-immediate, giving call edges.
    Spurious edges from misaligned decodes are possible; they only ADD reachability.
  * function starts = the set of call targets (+ roots); a site's enclosing function is the
    greatest function start <= site. This avoids the incomplete-recursive-descent bug that an
    earlier version of this script had (it silently missed 0x404962 -> 0x404fe8, which is the
    edge that matters here -- validated against disasm_fn before trusting any negative).
  * fn-level BFS from the roots; collect every call site in every reachable function.

Usage:
  python3 retaddr_reach.py --roots 0x402d28 --pivots 0x402f00,0x405000,0x406300,0x404d00
"""
import argparse, collections, bisect
from capstone import *
from capstone.arm import *

BASE = 0x400000
IMG = 'gcp_22.0.0.77_image_0x400000.bin'
CODE_LO, CODE_HI = 0x400000, 0x407000
CALLS = (ARM_INS_BL, ARM_INS_BLX)


def disasm_fn(img, md, start, limit=4000):
    """recursive descent within one function (kept for validation/diagnostics only)"""
    seen, work, calls = {}, [start], []
    while work:
        pc = work.pop()
        while True:
            if pc in seen or not (CODE_LO <= pc < CODE_HI):
                break
            got = list(md.disasm(img[pc - BASE:pc - BASE + 4], pc, count=1))
            if not got:
                break
            i = got[0]
            seen[pc] = i
            if len(seen) > limit:
                return seen, calls
            tgt = next((op.imm for op in i.operands if op.type == ARM_OP_IMM), None)
            if i.id in CALLS:
                calls.append((pc, pc + i.size, tgt))
                pc += i.size
                continue
            if i.id == ARM_INS_B:
                if tgt is not None and CODE_LO <= tgt < CODE_HI:
                    work.append(tgt)
                if i.cc in (ARM_CC_AL, ARM_CC_INVALID):
                    break
                pc += i.size
                continue
            if i.id in (ARM_INS_CBZ, ARM_INS_CBNZ):
                if tgt is not None:
                    work.append(tgt)
                pc += i.size
                continue
            if i.id == ARM_INS_BX or (i.id == ARM_INS_POP and 'pc' in i.op_str):
                break
            pc += i.size
    return seen, calls


def global_edges(img, md):
    """every BL/BLX-immediate found at any even address: (site, retaddr, target)"""
    out = []
    for pc in range(CODE_LO, CODE_HI, 2):
        got = list(md.disasm(img[pc - BASE:pc - BASE + 4], pc, count=1))
        if not got:
            continue
        i = got[0]
        if i.id in CALLS and i.size == 4:
            for op in i.operands:
                if op.type == ARM_OP_IMM and CODE_LO <= op.imm < CODE_HI:
                    out.append((pc, pc + i.size, op.imm))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--roots', default='0x402d28')
    ap.add_argument('--pivots', default='0x402f00,0x405000,0x406300,0x404d00')
    args = ap.parse_args()
    roots = [int(x, 16) for x in args.roots.split(',')]
    pivots = set(int(x, 16) for x in args.pivots.split(','))

    img = open(IMG, 'rb').read()
    md = Cs(CS_ARCH_ARM, CS_MODE_THUMB)
    md.detail = True

    edges = global_edges(img, md)
    fstarts = sorted(set([t for _, _, t in edges]) | set(roots))

    def encl(site):
        k = bisect.bisect_right(fstarts, site) - 1
        return fstarts[k] if k >= 0 else None

    fn_calls = collections.defaultdict(list)   # fn -> [(site, ret, tgt)]
    fn_graph = collections.defaultdict(set)
    for site, ret, tgt in edges:
        f = encl(site)
        if f is None:
            continue
        fn_calls[f].append((site, ret, tgt))
        fn_graph[f].add(tgt)

    # fn-level BFS
    reach, work = set(roots), list(roots)
    while work:
        f = work.pop()
        for g in fn_graph.get(f, ()):
            if g not in reach:
                reach.add(g)
                work.append(g)

    aligned = collections.defaultdict(list)
    for f in reach:
        for site, ret, tgt in fn_calls.get(f, ()):
            aligned[ret & ~0xFF].append((f, site, ret, tgt))

    print("call graph from %s: %d functions reachable, %d call sites, %d function starts known"
          % (args.roots, len(reach), sum(len(fn_calls.get(f, ())) for f in reach), len(fstarts)))
    print("\ndistinct ARM-mode entries reachable by LSB-zeroing a live saved LR: %d" % len(aligned))
    for a in sorted(aligned):
        mark = '  <== PIVOT' if a in pivots else ''
        print("  0x%06x   (%d call sites)%s" % (a, len(aligned[a]), mark))

    print("\n--- pivot intersection ---")
    hit = sorted(set(aligned) & pivots)
    if not hit:
        print("EMPTY: no call site on this chain returns into any pivot's 256-byte block.")
        print("Pivots wanted: %s" % ', '.join('0x%06x' % p for p in sorted(pivots)))
    for a in hit:
        print("0x%06x reachable via:" % a)
        for f, site, ret, tgt in aligned[a]:
            print("   fn 0x%06x  call@0x%06x -> 0x%06x   saved LR = 0x%06x" % (f, site, tgt, ret))


if __name__ == '__main__':
    main()