#!/usr/bin/env python3
"""Where does a TA's STACK live, and what is the TA address space's permission map?

Hardware probing cost one address per power cycle (a TA data abort wedges the sOS ring), so the
FE09 write-reach sweep could only find: shared buffer RW, image code/rodata RO, .bss RW, and
unmapped gaps. The stack and heap were never located.

The emulator settles it for free, because two things are real here:
  * `boot_sos_and_tas` boots sOS from the reset vector and the FIRMWARE'S OWN loader
    (`bl_load_ta_image`) loads+verifies the TAs, so the SP it installs is the real one.
  * The ARM MMU is real (FIDELITY-NOTES retraction: SCTLR.M=1, QEMU honours sOS's own
    short-descriptor tables). So walking TTBR0 gives the true mapping + permissions.

Method
  1. Boot + stage + run until modules are loaded and TA code executes.
  2. Range-filtered block hook over the TA window records the live SP (and CPSR mode) while
     TA code runs -> the stack VA, directly.
  3. Walk all TTBR0 L1 entries (TTBCR.N=4 => 256 entries, VA 0..0x10000000) and decode
     AP/APX/XN/domain for sections and small pages -> the full map, so we can say which
     regions are user-writable and where the stack/heap sit relative to the image.

Nothing here is a guess: SP is observed and permissions are decoded from the live tables.
"""
import sys, os, time, collections, json
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import unicorn
from unicorn import UC_HOOK_BLOCK, UC_HOOK_CODE
from unicorn.arm_const import UC_ARM_REG_SP, UC_ARM_REG_PC, UC_ARM_REG_R0, UC_ARM_REG_CPSR

import importlib.util
_s = importlib.util.spec_from_file_location(
    "tasmod", os.path.join(os.path.dirname(__file__), "boot_sos_and_tas.py"))
T = importlib.util.module_from_spec(_s); _s.loader.exec_module(T)
from pspemu import armmmu

TA_LO, TA_HI = 0x400000, 0x500000
MODES = {0x10: "USR", 0x11: "FIQ", 0x12: "IRQ", 0x13: "SVC", 0x17: "ABT", 0x1B: "UND", 0x1F: "SYS"}


def ap_str(ap, apx, xn):
    """ARMv7 short-descriptor access permissions -> (priv, user)."""
    table = {(0, 0): ("---", "---"), (0, 1): ("RW-", "---"), (0, 2): ("RW-", "R--"),
             (0, 3): ("RW-", "RW-"), (1, 0): ("---", "---"), (1, 1): ("R--", "---"),
             (1, 2): ("R--", "R--"), (1, 3): ("R--", "R--")}
    p, u = table.get((apx, ap), ("?", "?"))
    if xn:
        p = p[:2] + "-"; u = u[:2] + "-"
    else:
        p = p[:2] + "X"; u = u[:2] + ("X" if u != "---" else "-")
    return p, u


def walk(uc):
    """Full TTBR0 walk -> list of (va_lo, va_hi, kind, pa, priv, user, domain)."""
    base = armmmu.ttbr0(uc); n = armmmu.l1_entries(uc)
    rd = lambda a: int.from_bytes(uc.mem_read(a, 4), "little")
    out = []
    for i in range(n):
        e = rd(base + 4 * i)
        t = e & 3
        va = i << 20
        if t == 0:
            continue
        dom = (e >> 5) & 0xF
        if t == 2:                                   # 1 MB section
            ap = (e >> 10) & 3; apx = (e >> 15) & 1; xn = (e >> 4) & 1
            p, u = ap_str(ap, apx, xn)
            out.append((va, va + 0x100000, "section", e & 0xFFF00000, p, u, dom))
        elif t == 1:                                 # coarse -> L2, 256 small pages
            l2b = e & 0xFFFFFC00
            run = None
            for j in range(256):
                e2 = rd(l2b + 4 * j)
                pva = va + (j << 12)
                if e2 & 3:
                    ap = (e2 >> 4) & 3; apx = (e2 >> 9) & 1; xn = e2 & 1
                    p, u = ap_str(ap, apx, xn)
                    key = (p, u, dom)
                    if run and run[2] == key and run[1] == pva:
                        run = (run[0], pva + 0x1000, key, run[3])
                    else:
                        if run:
                            out.append((run[0], run[1], "pages", run[3], *run[2]))
                        run = (pva, pva + 0x1000, key, e2 & 0xFFFFF000)
                else:
                    if run:
                        out.append((run[0], run[1], "pages", run[3], *run[2]))
                        run = None
            if run:
                out.append((run[0], run[1], "pages", run[3], *run[2]))
    return out


def main():
    c, ccp, exc_log = T.build()
    print("staging images at PA %#x:" % T.IMAGE_PA)
    T.stage_all(c, verbose=False)

    loads = []; bs = set(); ta_hit = []
    # SP observations while TA code executes, bucketed by CPSR mode
    sp_by_mode = collections.defaultdict(list)
    sp_pc = []

    def code(uc, a, sz, u):
        if a in T.LOAD_RETS:
            loads.append((T.LOAD_RETS[a], uc.reg_read(UC_ARM_REG_R0)))

    def on_ta_block(uc, a, s, u):
        ta_hit.append(a)
        sp = uc.reg_read(UC_ARM_REG_SP)
        mode = MODES.get(uc.reg_read(UC_ARM_REG_CPSR) & 0x1F, "?")
        sp_by_mode[mode].append(sp)
        if len(sp_pc) < 40:
            sp_pc.append((a, sp, mode))

    c.uc.hook_add(UC_HOOK_BLOCK, on_ta_block, begin=TA_LO, end=TA_HI - 1)
    c.uc.hook_add(UC_HOOK_CODE, code, begin=0x2048E0, end=0x204B70)
    c.uc.hook_add(UC_HOOK_BLOCK, lambda uc, a, s, u: bs.add(a))
    c.uc.reg_write(UC_ARM_REG_SP, 0x3F000)

    def done():
        return sum(1 for _w, rv in loads if rv == 0) >= 4 and bool(ta_hit)

    t = time.time()
    T.run_system(c, seconds=float(os.environ.get("RUN_SECONDS", "60")),
                 done=None if os.environ.get("RUN_FULL") else done, verbose=True)
    print("  ran %.1fs, %d TA blocks, loads ok=%d"
          % (time.time() - t, len(ta_hit), sum(1 for _w, rv in loads if rv == 0)))

    print("\n=== TA live SP observations (the stack, directly) ===")
    if not sp_by_mode:
        print("  NO TA blocks executed -> nothing to report (negative control: do not infer)")
    for mode, sps in sorted(sp_by_mode.items()):
        print("  mode %-4s n=%-7d SP min=%#010x max=%#010x  (span %#x)"
              % (mode, len(sps), min(sps), max(sps), max(sps) - min(sps)))
    print("  first TA blocks (pc, sp, mode):")
    for a, sp, m in sp_pc[:12]:
        print("     pc=%#010x sp=%#010x %s" % (a, sp, m))

    uc = c.uc
    print("\n=== MMU state ===")
    print("  SCTLR=%#010x (M=%d)  TTBCR=%#x  TTBR0=%#010x  L1 entries=%d (VA 0..%#x)"
          % (armmmu.sctlr(uc), armmmu.mmu_enabled(uc), armmmu.ttbcr(uc),
             armmmu.ttbr0(uc), armmmu.l1_entries(uc), armmmu.l1_entries(uc) << 20))

    rows = walk(uc)
    print("\n=== full TTBR0 map (priv/user perms, X=exec) ===")
    print("  %-22s %-8s %-10s %-5s %-5s %s" % ("VA range", "kind", "PA", "priv", "user", "dom"))
    for lo, hi, kind, pa, p, u, dom in rows:
        tag = ""
        if lo <= 0x400000 < hi or (0x400000 <= lo < 0x500000): tag = " <- TA image window"
        if lo <= 0x4013000 < hi: tag = " <- SHARED BUFFER (params_VA)"
        for mode, sps in sp_by_mode.items():
            if any(lo <= s < hi for s in sps): tag += " <- *** TA STACK (%s) ***" % mode
        print("  %#010x-%#010x %-8s %#010x %-5s %-5s %-3d%s" % (lo, hi, kind, pa, p, u, dom, tag))

    # Which regions are USER-writable? that's what the FE09 store can reach.
    print("\n=== user-writable regions (what a TA-context store can hit) ===")
    for lo, hi, kind, pa, p, u, dom in rows:
        if u.startswith("RW"):
            print("  %#010x-%#010x  %-8s user=%s dom=%d" % (lo, hi, kind, u, dom))
    json.dump([[lo, hi, kind, pa, p, u, dom] for lo, hi, kind, pa, p, u, dom in rows],
              open("/tmp/ta_mmu_map.json", "w"), indent=1)
    print("\nwrote /tmp/ta_mmu_map.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())