#!/usr/bin/env python3
"""Test Addendum 146's reopening condition 1 directly, instead of asserting it.

SUPERSEDED / UNRELIABLE -- kept for the record. Its PTR heuristic flags 73 of 106 functions per
build, which is obviously wrong: it counts `ldr r0,[pc,#imm]` literal loads, arithmetic on status
values, and `pop {r4,...}` register RESTORES as pointer definitions. Addendum 147 replaces it with
the precise query (enumerate (frame, pivot) pairs whose pivot base is caller-saved, then read each
frame's epilogue by hand). Do not cite this script's output.

Add.146 closed the LSB-zero pivot line with a structural argument: the trick corrupts the saved LR
of a frame that must CONTAIN the host-aimed store in its subtree; such frames are command handlers;
handlers return status codes; so r0 at the pivot is always a small integer. The stated way to reopen
it was "a host-aimed store inside the subtree of a function that returns a POINTER", which Add.146
claimed does not exist in any build -- on reasoning, not measurement.

This script measures it. For every function in every build:
  * classify what it leaves in r0 at its return:
      STATUS  - every reaching definition is a small immediate (movs/movw/mov.w r0,#imm)
      PTR     - some reaching definition is a register that was produced by pointer arithmetic
                (add rN,rM,#imm / adds rN,rM,rK / ldr rN,[rM,...]) rather than an immediate
      UNKNOWN - no definition found before the terminator (r0 passes through from a callee)
  * compute its call subtree and look for register-offset stores (the host-aimed write shape)
  * report any PTR-returning function that has such a store in its subtree

A hit is a candidate only: store shape and pointer-ness are both heuristics and need hand checks.
An empty result makes Add.146's condition-1 claim measured rather than assumed.

Usage: python3 ptr_return_store.py
"""
import bisect, collections, glob, os, struct
from capstone import *
from capstone.arm import *

HERE = os.path.dirname(os.path.abspath(__file__))
VERS = os.path.join(HERE, '..', '..', 'psp_out', 'versions_win')
BASE = 0x400000
ST = (ARM_INS_STR, ARM_INS_STRB, ARM_INS_STRH, ARM_INS_STRD)
CALLS = (ARM_INS_BL, ARM_INS_BLX)
IMM_MOV = (ARM_INS_MOV, ARM_INS_MOVW, ARM_INS_MOVS)
PTR_PRODUCERS = (ARM_INS_ADD, ARM_INS_ADDW, ARM_INS_LDR, ARM_INS_SUB, ARM_INS_SUBS, ARM_INS_SUBW)


def load_build(path):
    body = open(path, 'rb').read()
    lo, hi = struct.unpack_from('<II', body, 4)
    return b'\x00' * 0x100 + body, lo, hi


def func_body(md, img, f, hi):
    """sequential decode from prologue to natural terminator"""
    out = []
    for insn in md.disasm(img[f - BASE:min(f + 0x600, hi) - BASE], f):
        out.append(insn)
        if insn.id == ARM_INS_POP and 'pc' in insn.op_str:
            break
        if insn.id == ARM_INS_BX and 'lr' in insn.op_str:
            break
    return out


def classify_r0(insns):
    """what does this function leave in r0 at return?"""
    ptr_regs = set()
    verdict = None
    for insn in insns:
        ops = insn.operands
        # track registers produced by pointer arithmetic
        if insn.id in PTR_PRODUCERS and ops and ops[0].type == ARM_OP_REG:
            srcs = [o for o in ops[1:]]
            if any(o.type == ARM_OP_REG for o in srcs) or any(o.type == ARM_OP_MEM for o in srcs):
                ptr_regs.add(ops[0].reg)
        # r0 definitions
        if ops and ops[0].type == ARM_OP_REG and ops[0].reg == ARM_REG_R0 and (ops[0].access & CS_AC_WRITE):
            if insn.id in IMM_MOV and len(ops) > 1 and ops[1].type == ARM_OP_IMM:
                verdict = 'STATUS' if verdict in (None, 'STATUS') else verdict
            elif insn.id in IMM_MOV and len(ops) > 1 and ops[1].type == ARM_OP_REG:
                verdict = 'PTR' if ops[1].reg in ptr_regs else (verdict or 'STATUS')
            elif insn.id in PTR_PRODUCERS:
                verdict = 'PTR'
    return verdict or 'UNKNOWN'


def main():
    for path in sorted(glob.glob(os.path.join(VERS, 'ta_22-0-*.bin'))):
        name = os.path.basename(path)
        img, lo, hi = load_build(path)
        mt = Cs(CS_ARCH_ARM, CS_MODE_THUMB); mt.detail = True
        edges = []
        for pc in range(lo, hi, 2):
            g = list(mt.disasm(img[pc - BASE:pc - BASE + 4], pc, count=1))
            if g and g[0].id in CALLS and g[0].size == 4:
                for op in g[0].operands:
                    if op.type == ARM_OP_IMM and lo <= op.imm < hi:
                        edges.append((pc, op.imm))
        fs = sorted(set(t for _, t in edges))

        def encl(s):
            k = bisect.bisect_right(fs, s) - 1
            return fs[k] if k >= 0 else None
        graph = collections.defaultdict(set)
        for s, t in edges:
            graph[encl(s)].add(t)

        kinds, stores = {}, {}
        for f in fs:
            body = func_body(mt, img, f, hi)
            kinds[f] = classify_r0(body)
            hit = []
            for insn in body:
                if insn.id in ST:
                    for op in insn.operands:
                        if op.type == ARM_OP_MEM and op.mem.index != 0:
                            hit.append((insn.address, "%s %s" % (insn.mnemonic, insn.op_str)))
            if hit:
                stores[f] = hit

        def subtree(root):
            seen, work = {root}, [root]
            while work:
                x = work.pop()
                for y in graph.get(x, ()):
                    if y not in seen:
                        seen.add(y); work.append(y)
            return seen

        nptr = sum(1 for v in kinds.values() if v == 'PTR')
        hits = []
        for f, k in kinds.items():
            if k != 'PTR':
                continue
            sub = subtree(f) - {f}
            for g2 in sub:
                if g2 in stores:
                    hits.append((f, g2, stores[g2][0]))
        print("%-28s fns=%-4d PTR-returning=%-4d store-fns=%-3d  PTR-with-store-in-subtree=%d"
              % (name, len(fs), nptr, len(stores), len(hits)))
        for f, g2, st in hits[:10]:
            print("    PTR fn 0x%06x  ->  store fn 0x%06x @0x%06x  %s" % (f, g2, st[0], st[1]))


if __name__ == '__main__':
    main()