#!/usr/bin/env python3
"""The whole system in one run: sOS + SYS_DRV + HDCP + DTM + RAS + ASD.

`boot_sos_and_tas.py` gets sOS, SYS_DRV, HDCP, DTM and RAS running but never
ASD -- ASD is driver-class and its boot-time staging slot is gated off (see
psp_out/FINDING-SOS-AND-TAS-BOOT-IN-EMULATOR.md), so on hardware it arrives
through the host ring instead. `load_asd.py` loads it that way but in isolation.
This does both in a single system, and then says which modules actually
EXECUTED rather than which ones loaded.

Attribution matters here. Every TA runs at VA 0x400000 in its OWN address
space, so a block count keyed on VA cannot distinguish HDCP from DTM from RAS
-- which is exactly what `boot_sos_and_tas`'s "TA/other=28 blocks" line does,
and why it can report TAs running without being able to name one. Blocks are
therefore attributed two ways that must agree:

  * address space -- ASD is based at 0xD00000, the TAs at 0x400000, SYS_DRV at
    0xE00000, sOS below 0x300000;
  * live physical page -- translated per block and matched against the staged
    image content, which separates HDCP/DTM/RAS inside the shared 0x400000
    window.
"""
import collections
import os
import struct
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.dirname(__file__))

import importlib.util                                            # noqa: E402
_b = importlib.util.spec_from_file_location(
    "bmmod", os.path.join(os.path.dirname(__file__), "boot_and_mailbox.py"))
BM = importlib.util.module_from_spec(_b); _b.loader.exec_module(BM)
M = BM.M

from pspemu import armmmu                                        # noqa: E402
from unicorn import UC_HOOK_BLOCK, UC_HOOK_CODE                  # noqa: E402
from unicorn.arm_const import UC_ARM_REG_R0                      # noqa: E402

# return sites of bl_load_ta_image inside bl_stage2_task (same as
# boot_sos_and_tas.LOAD_RETS) -- so this run shows the firmware's own loader
# accepting each module, not just code executing later.
LOAD_RETS = {0x2048E6: "slot0 driver (sys_drv)", 0x204A5C: "slot +0x6F8000 (asd)",
             0x204ABA: "slot +0x78C000 (ras)", 0x204B08: "slot +0x6C8000 (hdcp)",
             0x204B66: "slot +0x7A8000 (dtm)"}

PSP_OUT = BM.PSP_OUT
MODULES = ["drivers/sys_drv.bin", "hdcp_full.bin", "dtm_full.bin",
           "ras_full.bin", "asd_full.bin"]
SOS_HI = 0x300000
TA_LO, TA_HI = 0x400000, 0x500000          # HDCP / DTM / RAS, each in its own space
ASD_LO, ASD_HI = 0x00D00000, 0x00E00000
SD_LO, SD_HI = 0x00E00000, 0x00F00000
TA_IMAGE_PA = 0x04400000

# Liveness markers per module. "executing" is too weak a bar on its own -- DTM
# and RAS run only 8-9 blocks, which could be incidental. A TA is running in the
# meaningful sense when sOS has STARTED it: its entry trampoline runs, it
# executes its GP lifecycle entry, and it returns to the trampoline's SVC 0xF0
# message wait. The entry addresses are the ones decoded from each image's
# trampoline literal pool (see harness/dtm_ras.py), and the live system
# executing exactly those is an independent confirmation of that decode.
#
# ASD has no GP create phase -- its trampoline goes straight to the dispatcher
# (verified: *(0xD19884) == 0xD013B9) -- so its second marker is the dispatcher.
LIVENESS = {
    "hdcp_full.bin": (("trampoline", 0x4042A0), ("TA_CreateEntryPoint", 0x40139C)),
    "dtm_full.bin":  (("trampoline", 0x4011AC), ("TA_CreateEntryPoint", 0x401050)),
    "ras_full.bin":  (("trampoline", 0x4045E4), ("TA_CreateEntryPoint", 0x402B08)),
    "asd_full.bin":  (("trampoline", 0xD19868), ("dispatcher",          0xD013B8)),
}


def load_images():
    out = {}
    for n in MODULES:
        p = os.path.join(PSP_OUT, n)
        if os.path.exists(p):
            out[n] = open(p, "rb").read()
    return out


def attribute_page(c, pg, imgs):
    """Name the module a physical code page came from, by content.

    Several 48-byte slices are intersected: a single slice near the page start
    hits shared libc/prologue boilerplate and matches every image.
    """
    cand = set(imgs)
    for off in (0x40, 0x400, 0x800, 0xC00, 0x1000 - 0x40):
        try:
            sl = bytes(c.uc.mem_read(pg + off, 48))
        except Exception:
            continue
        if sl.count(0) > 40:
            continue
        cand &= {nm for nm, b in imgs.items() if sl in b}
        if len(cand) == 1:
            break
    return sorted(cand)


def main():
    imgs = load_images()
    c, ccp, irqs, exc = BM.build()
    run, _ = M.make_det_runner(c, irqs)

    hits = collections.Counter()

    def on_blk(u, a, s, x):
        hits[(a, armmmu.translate_live(u, a))] += 1

    for lo, hi in ((TA_LO, TA_HI), (ASD_LO, ASD_HI)):
        c.uc.hook_add(UC_HOOK_BLOCK, on_blk, begin=lo, end=hi - 1)
    sos_sd = collections.Counter()
    c.uc.hook_add(UC_HOOK_BLOCK, lambda u, a, s, x: sos_sd.__setitem__("sos", sos_sd["sos"] + 1),
                  begin=0x200000, end=SOS_HI - 1)
    c.uc.hook_add(UC_HOOK_BLOCK, lambda u, a, s, x: sos_sd.__setitem__("sys_drv", sos_sd["sys_drv"] + 1),
                  begin=SD_LO, end=SD_HI - 1)
    loads = []
    c.uc.hook_add(UC_HOOK_CODE,
                  lambda u, a, s_, x: loads.append((LOAD_RETS[a], u.reg_read(UC_ARM_REG_R0)))
                  if a in LOAD_RETS else None,
                  begin=0x2048E0, end=0x204B70)
    c.uc.ctl_flush_tb()

    # 1. boot sOS; bl_stage2_task loads and verifies the boot-time modules
    BM.drain_until(run,
                   lambda: int.from_bytes(c.uc.mem_read(M.RING_STATUS, 4), "little") == 0x80000000,
                   slices=12, blocks=500_000)
    ring_ready = int.from_bytes(c.uc.mem_read(M.RING_STATUS, 4), "little") == 0x80000000
    print("sOS booted, ring-ready word set : %s" % ring_ready)
    M.install_debug_gpio(c)

    # 2. bring the host command ring up
    box = M.Mailbox(c, irqs)
    c.uc.mem_write(BM.RING_PA, b"\x00" * BM.RING_SZ)
    rp = bytearray(0x1C)
    rp[0x10:0x14] = struct.pack("<I", BM.RING_PA)
    rp[0x18:0x1C] = struct.pack("<I", BM.RING_SZ)
    st, _raw = box.submit(2, bytes(rp), timeout=12, runner=run)
    print("ring init                       : status %s" % st)
    ring = M.Ring(c, irqs, BM.RING_PA, BM.RING_SZ, BM.CMDBUF_PA, run)

    # 3. ASD arrives over the ring, the way it does on hardware
    # SKIP_ASD=1 is a control, not a convenience: with the ring load skipped
    # ASD must show zero blocks and the run must report INCOMPLETE. If it still
    # claimed ASD was running, the attribution would be measuring something
    # other than the load.
    asd = None if os.environ.get("SKIP_ASD") else imgs.get("asd_full.bin")
    if asd:
        for base in (TA_IMAGE_PA, TA_IMAGE_PA + M.Ring.ALIAS):
            c.uc.mem_write(base, asd)
        ld = bytearray(0x20)
        ld[0:4] = struct.pack("<I", TA_IMAGE_PA)
        ld[8:12] = struct.pack("<I", len(asd))
        ring.submit(1, bytes(ld), timeout=10, probe=[])
        BM.drain_until(run, lambda: any(ASD_LO <= va < ASD_HI for va, _pa in hits),
                       slices=6, blocks=2_000_000)

    # Equal settle time on BOTH paths. Without this the SKIP_ASD control also
    # removed the run time the ASD drain happened to provide, so HDCP and DTM
    # stopped executing too and the control proved nothing about ASD. Give the
    # system the same budget either way, so the only variable is the ring load.
    for _ in range(6):
        run(2_000_000)

    # 4. attribute every executed block
    per_mod = collections.Counter()
    pages = collections.defaultdict(collections.Counter)
    vas = collections.defaultdict(set)
    for (va, pa), n in hits.items():
        pg = (pa & ~0xFFF) if pa else None
        if ASD_LO <= va < ASD_HI:
            per_mod["asd_full.bin"] += n
            pages["asd_full.bin"][pg] += n
            vas["asd_full.bin"].add(va)
        elif TA_LO <= va < TA_HI and pg is not None:
            who = attribute_page(c, pg, imgs)
            key = who[0] if len(who) == 1 else "unattributed"
            per_mod[key] += n
            pages[key][pg] += n
            vas[key].add(va)

    print("\nbl_load_ta_image results (the firmware's own loader + signature check):")
    for what, rv in loads:
        print("   %-26s rv=%#010x %s" % (what, rv, "LOADED" if rv == 0 else "rejected"))
    if not loads:
        print("   (none observed)")

    print("\nexecution by module (VA space + live physical page):")
    print("   %-22s %-10s %s" % ("module", "blocks", "physical code pages"))
    for nm in ("drivers/sys_drv.bin", "hdcp_full.bin", "dtm_full.bin",
               "ras_full.bin", "asd_full.bin"):
        n = per_mod.get(nm, 0)
        pg = [hex(p) for p in pages[nm] if p is not None]
        if nm == "drivers/sys_drv.bin":
            n = sos_sd["sys_drv"]; pg = ["(VA 0xE00000 space)"] if n else []
        print("   %-22s %-10d %s" % (nm, n, ", ".join(pg) or "-"))
    print("   %-22s %-10d" % ("sOS", sos_sd["sos"]))
    if per_mod.get("unattributed"):
        print("   (%d blocks on pages that matched more than one image)"
              % per_mod["unattributed"])

    n_loaded = sum(1 for _w, rv in loads if rv == 0)
    print("\nliveness (started by sOS, not merely touched):")
    live = {}
    for nm, marks in LIVENESS.items():
        got = []
        for label, addr in marks:
            hit = any(addr <= v < addr + 0x60 for v in vas[nm])
            got.append("%s %s" % (label, "yes" if hit else "NO"))
        live[nm] = all("NO" not in g for g in got)
        print("   %-16s %s" % (nm, "  |  ".join(got)))

    running = {
        "sOS": sos_sd["sos"] > 0,
        "SYS_DRV": sos_sd["sys_drv"] > 0,
        "HDCP": per_mod.get("hdcp_full.bin", 0) > 0 and live["hdcp_full.bin"],
        "DTM": per_mod.get("dtm_full.bin", 0) > 0 and live["dtm_full.bin"],
        "RAS": per_mod.get("ras_full.bin", 0) > 0 and live["ras_full.bin"],
        "ASD": per_mod.get("asd_full.bin", 0) > 0 and live["asd_full.bin"],
    }
    print("\n== summary ==")
    for k, v in running.items():
        print("   %-8s executing: %s" % (k, "yes" if v else "NO"))
    # ASD is deliberately absent from the boot-time list: its staging slot is
    # gated off (bl_stage2_task branches past it), which is why it comes over
    # the ring. Its evidence is that ASD code executes at all, which requires
    # tee_load_ta to have verified it -- load_asd.py pins that down further with
    # bit-flip negative controls.
    print("   modules loaded+verified at boot by the firmware itself: %d" % n_loaded)
    print("   ASD loaded over the host ring (boot slot is gated off)  : %s"
          % ("yes" if running["ASD"] else "no"))
    print("   faults: %d   ccp: %s" % (len(c.faults), dict(ccp.stats())))

    if all(running.values()):
        print("\n   RESULT: full system running -- sOS + SYS_DRV + all four TAs,")
        print("   each attributed by address space AND live physical page, and each")
        print("   started by sOS (trampoline + lifecycle entry executed).")
        return 0
    print("\n   RESULT: INCOMPLETE -- %s not executing"
          % ", ".join(k for k, v in running.items() if not v))
    return 1


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