#!/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


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 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")
    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. Nothing is required
        # of the link from here on. The verdict is *issued* before the SPM, but
        # the erase takes the code that would have finished saying it, and how
        # much of it survives is the chip's business - an erase removes one page
        # and nothing else, so a loader whose command loop lives past the page
        # erased will prompt as usual where one with 128-byte pages goes with
        # the stub. None of that is the claim. The claim is that the erase
        # reached flash, and the dump is both the only witness for it and a
        # better one: it tells "accepted and performed" from "merely answered".
        port.write(frame)
        try:
            port.read_exact(2, 2.0)
        except pb.Error:
            pass
        port.close()
    finally:
        device.stop()

    # Ground truth: the simulator's flash, not the loader's opinion of it.
    flash = open(dump, "rb").read()
    if flash[base : base + page] != b"\xff" * page:
        fail("the sealed erase did not reach flash - the running page is intact")
    print("pbselfwrite: the running slot is refused unsealed and erased sealed")


main()