#!/usr/bin/env python3
"""Boot sOS, then run the stage-2 payload task -- the sOS-side component loader.

`bl_main` spawns exactly one task via `SVC 0x51` (entry VA 0x204895). IDA had no
function defined there at all, because nothing statically references it; it is
reachable only through the task spawn. It is named `bl_stage2_task` in the IDB
now, and it runs:

    bl_stage2_init -> bl_init_obj_slots -> bl_load_ta_image

i.e. it is the sOS firmware-authenticity path. Reaching it required fixing four
separate things, each of which failed silently:

1. **The MMU is ON.** sOS sets SCTLR.M during boot (SCTLR=0xC5107D,
   TTBR0=0xFC48, TTBCR=0x24). Unicorn is QEMU and implements ARMv7 translation
   faithfully, so after boot every access goes through sOS's page tables and
   `uc.mem_map()` only decides what physical memory exists behind them.

2. **`SVC 0x79` (VA->PA translate) cannot be stubbed to 0.** `bl_main` calls it
   to obtain the physical address of its own code, then passes that to
   `SVC 0x62` to re-map it. With the stub returning 0, bl_main re-mapped
   VA 0x200000 to *physical 0* and destroyed the CODE2 mapping: VA 0x203b5a
   went from PA 0x15B5A (correct, = VA-0x1EE000) to PA 0x3B5A, after which the
   task executed unrelated bytes. Nothing reported an error.

3. **`SVC 0x62` (map pages) must really edit the page tables.** Returning 0
   claims success while creating no translation, and the VA faults later.

4. **CPU mode.** Boot ends in User mode (CPSR=0x30). Kernel pages are
   privileged-only, so calling kernel code from User mode takes a permission
   fault.

And the trap that made all of the above expensive to diagnose:

   **Once the MMU is on, `count=` does NOT bound emu_start.** A faulting access
   retires zero instructions, so the counter never decrements and the abort
   loop runs forever -- `count=1000` hangs exactly as long as `count=200000`.
   Use `timeout=` (wall clock) and `armmmu.require_mapped()` instead.

Current state: the task reaches `bl_load_ta_image` and then stalls inside
`bl_memcpy`, because no TA image is staged at the physical address it reads
(from 0x2060E8+0x38). Staging a real signed container there is the next step.
"""
import sys, os, time, json
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import unicorn
from unicorn import UC_PROT_ALL, UC_HOOK_BLOCK
from unicorn.arm_const import (UC_ARM_REG_SP, UC_ARM_REG_PC, UC_ARM_REG_LR,
                               UC_ARM_REG_CPSR, UC_ARM_REG_R0, UC_ARM_REG_R1,
                               UC_ARM_REG_R2)

import importlib.util
_s = importlib.util.spec_from_file_location(
    "bootmod", os.path.join(os.path.dirname(__file__), "boot_sos_full.py"))
B = importlib.util.module_from_spec(_s); _s.loader.exec_module(B)

from pspemu.ccp import CCP, CCP_MMIO_SIZE
from pspemu import armmmu
from pspemu.memmap import MAGIC_RET

STAGE2_TASK = 0x204894          # bl_stage2_task (thumb)
BL_MAIN = 0x203A30
CODE2_BIAS = 0x1EE000


def build(verbose=True):
    c, img = B.load_sos()
    for s in range(0x40, 0xC0):
        c.svc.register(s, lambda comp, n: 0)
    for a in B.CP15_SKIP:
        c.write(a, B.THUMB_NOP2)
    for base, size in B.AXI_MAPS_CCP:
        try:
            c.uc.mem_map(base, size, UC_PROT_ALL)
        except unicorn.UcError:
            pass
    ccp = CCP(c, base=B.CCP_WINDOW_VA, size=CCP_MMIO_SIZE)

    c.uc.reg_write(UC_ARM_REG_SP, 0x3F000)
    c.uc.emu_start(0x98, 0xFF0000, count=100000)      # -> bl_main idle loop
    assert 0x203B5A <= c.uc.reg_read(UC_ARM_REG_PC) < 0x203B70, "boot did not idle"

    pt = armmmu.PageTableEditor(c.uc)
    assert pt.selftest(0x203B5A), "page-table walk selftest failed"

    def svc79(comp, n):                               # VA -> PA
        _kind, pa = pt.describe(comp.uc.reg_read(UC_ARM_REG_R0))
        return pa or 0

    def svc62(comp, n):                               # map pages
        va = comp.uc.reg_read(UC_ARM_REG_R0)
        pa = comp.uc.reg_read(UC_ARM_REG_R1)
        np = comp.uc.reg_read(UC_ARM_REG_R2)
        if np == 0 or np > 0x4000 or pa == 0:
            return 0
        pt.ensure_backed(pa, np * 0x1000)
        try:
            pt.map_range(va, pa, np)
            return 0
        except Exception:
            return 1

    c.svc.register(0x79, svc79)
    c.svc.register(0x62, svc62)

    # Re-run bl_main so its mapping SVCs take effect for real this time.
    c.uc.reg_write(UC_ARM_REG_SP, 0x3F000)
    c.uc.emu_start(BL_MAIN | 1, 0xFF0000, count=400000)

    # INVARIANT: bl_main's own re-mapping must not have moved its code.
    kind, pa = pt.describe(0x203B5A)
    assert pa == 0x203B5A - CODE2_BIAS, (
        "CODE2 mapping corrupted: VA 0x203b5a -> PA %#x, expected %#x "
        "(this is what a stubbed SVC 0x79 does)" % (pa or 0, 0x203B5A - CODE2_BIAS))

    pt.ensure_backed(MAGIC_RET, 0x1000)
    pt.map_page(MAGIC_RET, MAGIC_RET)
    if verbose:
        print("boot OK; CODE2 mapping intact (VA 0x203b5a -> PA %#x)" % pa)
    return c, ccp, pt


def stage_image(c, pt, path, pa=0x00A00000):
    """Stage a real signed $PS1 container where bl_stage2_task looks for it.

    The task loads a DRIVER image: `sys_drv.bin` carries 'DR' (0x5244) at
    body+0x14, with body+0x0c = 0xEB6020 and body+0x18 = 0x27000, matching
    SYS_DRV's known layout. So stage 2 is loading SYS_DRV.

    Two addresses matter:
      * the container must be visible at VA 0x20A000 (its 0x100-byte signed
        header) / 0x20A100 (the TA header), where bl_load_ta_image reads it;
      * its physical address must be published at VA 0x2060E8+0x38, which
        bl_stage2_task reads with `LDRD R1,R0,[R9,#0x38]`.

    NOTE: `uc.mem_read/mem_write` are PHYSICAL and bypass the MMU. Writing a
    guest global by its VA silently lands somewhere else -- VA 0x206120 is
    PA 0x18120. Always translate first.
    """
    img = open(path, "rb").read()
    pt.ensure_backed(pa, len(img) + 0x2000)
    c.uc.mem_write(pa, img)
    pt.map_range(0x20A000, pa, (len(img) + 0xFFF) // 0x1000)

    def vwr(va, val):
        _k, p = pt.describe(va)
        c.uc.mem_write(p, int(val).to_bytes(4, "little"))

    vwr(0x206120, pa)       # pa_lo
    vwr(0x206124, 0)        # pa_hi
    return img


def install_svc64(c, pt, slot=0):
    """SVC 0x64 = create process. An OUT-PARAMETER call: it must write a slot
    index to *R1. Left as a stub it returns 0 ("success") while the caller's
    pre-set 0xFF remains, and bl_load_ta_image rejects with 0x25 because
    `ta_slot >= 0x20`.

    WARNING: this minimal model is NOT sufficient. It gets past the slot check,
    but real SVC 0x64 also builds the new process's address space and PCB;
    without that the loader spins forever re-clearing PCB state (24M blocks /
    120s with only 70 distinct blocks -- an infinite loop, not slow progress).
    Faithful process creation is the next piece of work.
    """
    from unicorn.arm_const import UC_ARM_REG_R1

    def svc64(comp, n):
        _k, pa = pt.describe(comp.uc.reg_read(UC_ARM_REG_R1))
        if pa is not None:
            comp.uc.mem_write(pa, bytes([slot]))
        return 0
    c.svc.register(0x64, svc64)


def run_task(c, pt, seconds=8):
    pt.require_mapped(pc=STAGE2_TASK, sp=0x287928, ret=MAGIC_RET)
    blocks = []
    c.uc.hook_add(UC_HOOK_BLOCK, lambda uc, a, s, u: blocks.append(a))
    c.uc.reg_write(UC_ARM_REG_SP, 0x287928)
    c.uc.reg_write(UC_ARM_REG_LR, MAGIC_RET | 1)
    t0 = time.time()
    try:
        # timeout=, NOT count= : count does not bound a fault loop under the MMU
        c.uc.emu_start(STAGE2_TASK | 1, MAGIC_RET, timeout=seconds * 1_000_000)
    except Exception as e:
        return blocks, str(e), time.time() - t0
    return blocks, None, time.time() - t0


def diagnose_stall(c, pt, blocks):
    """Distinguish "still working" from "wedged in an unbounded MMU abort loop".

    FIDELITY-NOTES describes the trap: an access to a VA with no guest
    translation retires ZERO instructions, so the faulting block restarts
    forever. It presents as "the emulator is slow", and the diagnostic pair is
      * `c.faults` EMPTY -- UC_HOOK_MEM_INVALID never fires, because the HOST
        memory is mapped; it is the GUEST translation that is missing; and
      * a tiny number of distinct blocks executing enormously many times.

    This harness used to report that state as the bland "timeout (still
    running)", which reads like progress. Measured here: 571k blocks, 32
    distinct, PC pinned at 0x200170 inside the ARM memcpy -- and the memcpy was
    entered exactly ONCE (len 0x20, from LR 0x204061), so it is not a long copy,
    it is one copy that can never complete.

    Returns a dict describing the stall, or None if the run looks live.
    """
    from unicorn.arm_const import (UC_ARM_REG_PC, UC_ARM_REG_R0, UC_ARM_REG_R1,
                                   UC_ARM_REG_R2, UC_ARM_REG_R3, UC_ARM_REG_R4,
                                   UC_ARM_REG_R5, UC_ARM_REG_R6, UC_ARM_REG_R7,
                                   UC_ARM_REG_SP, UC_ARM_REG_LR)
    if not blocks:
        return None
    distinct = len(set(blocks))
    # A live run keeps discovering blocks; a wedged one cycles a handful.
    if distinct > 64 or len(blocks) < 10000 or c.faults:
        return None
    pc = c.uc.reg_read(UC_ARM_REG_PC)
    regs = {"PC": pc, "SP": c.uc.reg_read(UC_ARM_REG_SP),
            "LR": c.uc.reg_read(UC_ARM_REG_LR)}
    for i, r in enumerate((UC_ARM_REG_R0, UC_ARM_REG_R1, UC_ARM_REG_R2,
                           UC_ARM_REG_R3, UC_ARM_REG_R4, UC_ARM_REG_R5,
                           UC_ARM_REG_R6, UC_ARM_REG_R7)):
        regs["R%d" % i] = c.uc.reg_read(r)
    unmapped = []
    for name, va in regs.items():
        if va < 0x1000:
            continue
        try:
            kind, pa = pt.describe(va)
        except Exception:
            continue
        if pa is None:
            unmapped.append((name, va, kind))
    # Only call it an ABORT LOOP when an unmapped VA is actually present.
    # Measured counter-example: after hand-mapping VA 0x20A000 the task moves on
    # and then spins at 0x205DC0 with 51 distinct blocks and NO unmapped
    # register -- a busy-wait, not a fault loop. Block RATE separates them too:
    # the abort loop retires zero instructions (542k blocks in 8s) while the
    # busy-wait runs flat out (15.1M in 8s). Reporting both as "abort loop"
    # would be a confident wrong diagnosis.
    kind = "abort_loop" if unmapped else "spin"
    return dict(pc=pc, blocks=len(blocks), distinct=distinct,
                unmapped=unmapped, kind=kind)


def main():
    c, ccp, pt = build()
    blocks, err, dt = run_task(c, pt)
    print("bl_stage2_task: %d blocks, %d distinct, pc=%#x, %.1fs"
          % (len(blocks), len(set(blocks)), c.uc.reg_read(UC_ARM_REG_PC), dt))
    stall = diagnose_stall(c, pt, blocks)
    if stall is None:
        print("  stop reason: %s" % (err or "timeout (still running)"))
    else:
        if stall["kind"] == "abort_loop":
            print("  stop reason: STALLED -- unbounded MMU abort loop, not progress")
        else:
            print("  stop reason: spinning (busy-wait), not an abort loop")
        print("     PC pinned at %#x, %d blocks but only %d distinct, c.faults empty"
              % (stall["pc"], stall["blocks"], stall["distinct"]))
        if stall["unmapped"]:
            for name, va, kind in stall["unmapped"]:
                print("     %-3s = %#010x has NO guest translation (%s)"
                      % (name, va, kind))
            print("     -> the faulting access can never retire, so the block")
            print("        restarts forever. Map that VA (or model the SVC that")
            print("        should have mapped it) rather than raising the timeout.")
        else:
            print("     no register holds an unmapped VA, so this is NOT a"
                  " translation fault -- it is a genuine busy-wait")
    print("  ccp: %s errors=%d" % (dict(ccp.stats()), len(ccp.errors)))
    json.dump(sorted(set(blocks)), open("/tmp/task_blocks.json", "w"))
    # success == we got into the loader, not merely that something ran
    reached_loader = any(0x203F24 <= b < 0x204100 for b in blocks)   # bl_load_ta_image
    print("  reached bl_load_ta_image: %s" % reached_loader)
    return 0 if reached_loader else 1


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