#!/usr/bin/env python3
"""Derive the loader's TA memory-layout RULE: where do code/data/stack/heap land?

`ta_stack_and_map.py` showed one TA's live SP (0x40c000) and that the stack sits just above the
image behind a guard page. To PREDICT another TA's stack (e.g. GCP, which this boot does not
load) the rule must be derived, not guessed -- so group observations by process.

Each TA process has its own L1 table, so TTBR0 identifies the process. For every TA basic block
we record (TTBR0, SP). Then each distinct TTBR0's table is walked independently, giving that
process's regions + permissions, and the stack is identified as the run containing its observed
SP values. With four TAs of different image sizes (hdcp/asd/ras/dtm) the layout arithmetic is
determined rather than inferred from a single instance.
"""
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_R0, UC_ARM_REG_CPSR, UC_ARM_REG_CP_REG

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


def ap_str(ap, apx, xn):
    table = {(0, 0): ("---", "---"), (0, 1): ("RW-", "---"), (0, 2): ("RW-", "R--"),
             (0, 3): ("RW-", "RW-"), (1, 1): ("R--", "---"), (1, 2): ("R--", "R--"),
             (1, 3): ("R--", "R--"), (1, 0): ("---", "---")}
    p, u = table.get((apx, ap), ("?", "?"))
    p = p[:2] + ("-" if xn else "X")
    u = u[:2] + ("-" if (xn or u == "---") else "X")
    return p, u


def walk_base(uc, base, n):
    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:
            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:
            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()
    T.stage_all(c, verbose=False)
    loads = []; ta_hit = []
    obs = collections.defaultdict(lambda: collections.Counter())   # ttbr0 -> Counter(sp)
    pcs = collections.defaultdict(set)

    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(uc, a, s, u):
        ta_hit.append(a)
        if (uc.reg_read(UC_ARM_REG_CPSR) & 0x1F) != 0x10:
            return                                   # USR mode only = the TA itself
        t0 = uc.reg_read(UC_ARM_REG_CP_REG, (15, 0, 0, 2, 0, 0, 0, 0))
        obs[t0][uc.reg_read(UC_ARM_REG_SP)] += 1
        pcs[t0].add(a)

    c.uc.hook_add(UC_HOOK_BLOCK, on_ta, begin=TA_LO, end=TA_HI - 1)
    c.uc.hook_add(UC_HOOK_CODE, code, begin=0x2048E0, end=0x204B70)
    c.uc.reg_write(UC_ARM_REG_SP, 0x3F000)

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

    T.run_system(c, seconds=float(os.environ.get("RUN_SECONDS", "90")),
                 done=None if os.environ.get("RUN_FULL") else done, verbose=True)
    print("TA blocks=%d  loads ok=%d  distinct TA processes (TTBR0)=%d"
          % (len(ta_hit), sum(1 for _w, rv in loads if rv == 0), len(obs)))

    n = armmmu.l1_entries(c.uc)
    report = {}
    for t0, spc in sorted(obs.items()):
        base = t0 & ~((1 << (14 - (armmmu.ttbcr(c.uc) & 7))) - 1)
        rows = [r for r in walk_base(c.uc, base, n) if TA_LO <= r[0] < TA_HI]
        sps = sorted(spc)
        print("\n=== TA process TTBR0=%#010x (%d blocks, %d distinct SP) ===" % (t0, sum(spc.values()), len(sps)))
        print("    SP observed: %s%s" % (', '.join(hex(s) for s in sps[:6]),
                                         ' ...' if len(sps) > 6 else ''))
        print("    %-24s %-7s %-5s %-5s %s" % ("VA range", "kind", "priv", "user", "note"))
        stack = None
        for lo, hi, kind, pa, p, u, dom in rows:
            note = ""
            if any(lo <= s <= hi for s in sps):
                note = "<== STACK (top %#x, size %#x)" % (hi, hi - lo); stack = (lo, hi)
            if p.endswith("X"):
                note = note or "code"
            print("    %#010x-%#010x %-7s %-5s %-5s %s" % (lo, hi, kind, p, u, note))
        # image extent = lowest exec page .. highest RW page below the stack
        execs = [r for r in rows if r[4].endswith("X")]
        img_lo = min(r[0] for r in execs) if execs else None
        pre = [r for r in rows if stack and r[1] <= stack[0]]
        img_hi = max(r[1] for r in pre) if pre else None
        if stack and img_hi:
            print("    RULE: image/data end %#x -> guard %#x -> stack [%#x,%#x) size %#x"
                  % (img_hi, stack[0] - img_hi, stack[0], stack[1], stack[1] - stack[0]))
            report[hex(t0)] = dict(img_lo=hex(img_lo), data_end=hex(img_hi),
                                   guard=hex(stack[0] - img_hi), stack_lo=hex(stack[0]),
                                   stack_hi=hex(stack[1]), stack_size=hex(stack[1] - stack[0]))
    json.dump(report, open("/tmp/ta_stack_rule.json", "w"), indent=1)
    print("\nwrote /tmp/ta_stack_rule.json")
    return 0


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