#!/usr/bin/env python3
"""The seal, gated on the one command that proves it: erase the page the
loader is executing from.

pureboot 9 dropped the running-slot write guard, so this command is now
permitted - that is what lets a resident copy plant something in its own slot,
which on a chip whose boot section *is* the loader slot is the only route a
self-update has. Permitted means the loader must actually do it, and the only
honest proof is the flash afterwards.

What stands in the guard's place is the seal, and the two halves are tested
against each other here: the identical destructive command, refused when its
seal is wrong and honoured when it is right. A test that only showed the
refusal would pass just as well against a loader that ignores SPM entirely.

Usage: pbselfwrite.py <device_bin> <pureboot_elf> <mcu> <hz> <base_hex> <page>
                      <baud> <tool_py> <workdir>
"""

import os
import sys
import time


# How long the device is left running after the sealed command, for the frame
# to arrive at the wire's rate and the erase to reach flash. It is spent in
# full on every attempt, because the only cheaper signal - the loader answering
# - is the one that precedes the SPM.
settle_seconds = 2.0
settle_attempts = 3


def fail(message):
    print(f"FAIL: {message}")
    sys.exit(1)


def sealed_frame(pb, op, space, address, count):
    """A command header and its seal, built here rather than borrowed from the
    tool: this test is about what the loader accepts, and a probe that shares
    the host's frame builder cannot tell a wrong frame from a wrong loader."""
    head = bytes((op, pb.selector(space, address), address & 0xFF,
                  (address >> 8) & 0xFF, count & 0xFF))
    seal = pb.SEAL
    for byte in head:
        seal ^= byte
    return head + bytes((seal,))


def scenario(pb, pbsim, device_bin, elf, mcu, hz, base_hex, page, baud, dump, base, settle):
    """The whole scenario once, answering with the running page as the
    simulator's own flash holds it afterwards."""
    device = pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, dump)
    try:
        port = pb.Port(device.pty, baud)
        loader = pb.Loader(port)
        info = loader.connect(25)
        if info.version < pb.SEALED_LOADER:
            fail(f"this test is for pureboot {pb.SEALED_LOADER} and later, not {info.version}")

        # Erase the first page of the running slot: the entry stub and the
        # command loop are both in it, so a loader that performs this does not
        # answer again. Nothing else in the protocol is as sharp a probe.
        frame = sealed_frame(pb, pb.OP_WRITE, pb.SP_SPM, base, pb.SPM_ERASE)

        # Red: the same command with one bit wrong in its seal. Refused before
        # anything happens, and the loader is still there to say so.
        broken = bytearray(frame)
        broken[-1] ^= 0x01
        port.write(bytes(broken))
        if port.read_exact(1, 5.0) != pb.NAK:
            fail("an unsealed erase of the running page was not refused")
        if port.read_exact(1, 5.0) != pb.PROMPT:
            fail("the loader did not re-prompt after refusing the erase")
        alive = loader.read_flash(base, 8)
        if alive == b"\xff" * 8:
            fail("the refused erase happened anyway - the running page reads erased")

        # Green: the identical command, correctly sealed. The claim is that the
        # erase reached flash, and the dump is the only witness for it - it
        # tells "accepted and performed" from "merely answered".
        #
        # Nothing the link says may shorten the settle, because the verdict is
        # issued *before* the SPM: a prompt reply means the erase has not
        # happened yet, and stopping the device on it races the write. The link
        # dying here is an expected outcome rather than a failure - the loader
        # is erasing its own command loop - so a closing pty ends the settle
        # instead of escaping it.
        port.write(frame)
        deadline = time.monotonic() + settle
        while time.monotonic() < deadline:
            try:
                port.read_available(0.1)
            except (pb.Error, OSError):
                break
        try:
            port.close()
        except OSError:
            pass
    finally:
        device.stop()

    # Ground truth: the simulator's flash, not the loader's opinion of it.
    return open(dump, "rb").read()[base : base + page]


def main():
    device_bin, elf, mcu, hz, base_hex, page, baud, tool, workdir = sys.argv[1:]
    base, page, baud = int(base_hex, 0), int(page), int(baud)
    sys.path.insert(0, os.path.dirname(os.path.abspath(tool)))
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    import pbsim
    import pureboot as pb

    os.makedirs(workdir, exist_ok=True)
    dump = os.path.join(workdir, "dump.bin")

    # A settle is wall clock and the erase is the simulator's progress through
    # it, so a loaded machine needs more of the first for the same amount of
    # the second - which is what made a single fixed wait flaky under the gate's
    # own parallelism. Doubling until the claim holds keeps the claim exact: a
    # loader that does not erase fails every attempt, and only the budget moves.
    settle = settle_seconds
    for _ in range(settle_attempts):
        if scenario(pb, pbsim, device_bin, elf, mcu, hz, base_hex, page, baud, dump, base, settle) == b"\xff" * page:
            print("pbselfwrite: the running slot is refused unsealed and erased sealed")
            return
        settle *= 2
    fail("the sealed erase did not reach flash - the running page is intact")


main()