#!/usr/bin/env python3
"""A link that damages bytes, deliberately and reproducibly, against the seal
that exists for it.

The board this was written for loses and mangles bytes on its own serial path,
and the failure that made it matter — a page-fill byte lost, the stream one
byte out, a page-address byte arriving where an SPMCSR value belongs — is not
reachable by asking a healthy link nicely. So the damage is injected here, at
a named byte index rather than a probability: a failing case is a case that
fails again.

Every check is a pair. The same bit flipped in the same field is applied on
one side of the seal and then the other: *after* the host seals the header,
which is a mangled command and must be refused, and *before*, which is a
well-formed command for something else and must be obeyed. Only the pair
proves anything — a test that showed the refusal alone would pass against a
loader that had simply stopped doing SPM, and one that showed the corruption
alone would not say what caught it.

Usage: pbglitch.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 frame(pb, op, space, address, count, damage=None, before_seal=False):
    """A sealed command header, optionally with one byte damaged.

    `damage` is (index, mask). Applied before the seal is computed it produces
    a valid command for whatever the damaged fields now say; applied after, a
    command whose seal no longer matches its own body — which is the shape a
    link fault actually has."""
    head = bytearray((op, pb.selector(space, address), address & 0xFF,
                      (address >> 8) & 0xFF, count & 0xFF))
    if damage and before_seal:
        head[damage[0]] ^= damage[1]
    seal = pb.SEAL
    for byte in head:
        seal ^= byte
    out = bytearray(head + bytes((seal,)))
    if damage and not before_seal:
        out[damage[0]] ^= damage[1]
    return bytes(out)


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}")

        # A page of known bytes to watch. Everything below aims at it, so
        # "nothing happened" is a readable claim rather than an absence.
        marker = bytes((0x40 + (i & 0x3F)) for i in range(page))
        loader.write_page(0, marker)
        if loader.read_flash(0, page) != marker:
            fail("the marker page did not survive an undamaged write")

        # Every field of the header, one bit each. A damaged seal must be
        # refused, the loader must re-prompt, and the page must be untouched —
        # and it is the erase being aimed at it, so a single escape is visible.
        for index in range(6):
            bad = frame(pb, pb.OP_WRITE, pb.SP_SPM, 0, pb.SPM_ERASE, damage=(index, 0x01))
            port.write(bad)
            answer = port.read_exact(1, 5.0)
            if answer != pb.NAK:
                fail(f"a header damaged in byte {index} was not refused (got {answer.hex()})")
            if port.read_exact(1, 5.0) != pb.PROMPT:
                fail(f"no prompt after refusing a header damaged in byte {index}")
            if loader.read_flash(0, page) != marker:
                fail(f"damage in byte {index} reached flash — the marker page changed")

        # The fill, whose payload is the protocol's one unacked burst: a
        # refused fill must be refused *before* the page is sent, or the host
        # is left pushing 128 bytes into a loader reading commands. Nothing is
        # sent after the verdict here, and the very next command must be
        # understood — that is the whole claim.
        bad = frame(pb, pb.OP_FILL, pb.SP_FLASH, 0, page, damage=(3, 0x80))
        port.write(bad)
        if port.read_exact(1, 5.0) != pb.NAK:
            fail("a damaged fill header was not refused")
        if port.read_exact(1, 5.0) != pb.PROMPT:
            fail("no prompt after refusing a damaged fill header")
        if loader.identity().raw != info.raw:
            fail("the loader was out of step after refusing a fill")

        # The pair's other half. The identical flip, applied before the seal:
        # a well-formed erase of the page one bit away from the one intended.
        # It must be obeyed — otherwise the refusals above prove nothing about
        # the seal and only that this loader stopped erasing.
        port.write(frame(pb, pb.OP_WRITE, pb.SP_SPM, 0, pb.SPM_ERASE,
                         damage=(2, 0x01), before_seal=True))
        if port.read_exact(1, 5.0) != pb.PROMPT:
            fail("a correctly sealed erase was refused")
        if port.read_exact(1, 5.0) != pb.PROMPT:
            fail("no prompt after a correctly sealed erase")
        if loader.read_flash(0, page) != b"\xff" * page:
            fail("the sealed erase did not reach flash — the marker page is intact")
        port.close()
    finally:
        device.stop()
    print("pbglitch: damaged headers are refused, the identical damage sealed is obeyed")


main()