#!/usr/bin/env python3
"""The reset walk region, end to end - the brick this refusal exists for.

BOOTRST programmed on a boot-sectioned mega enters reset at the boot section's
base, and the loader is reached only by walking up across erased flash. An
application that grew into that span is executed by the reset instead: it
happened to an ssd1306 board on a 328P with hfuse 0xdc, and with no reset edge
brought out, an ICE was the only way back.

The plain --flash is the path under test. It asks for no fuses of its own,
which is exactly why it once skipped the check and wrote the image anyway.

The runner still resets into the loader - the walk itself is not modeled. What
is modeled is the fuse read, so the refusal is decided from the bytes silicon
would have answered (-f, test/pureboot_device.cpp).

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

import os
import sys


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


def intel_hex(chunks):
    """Intel HEX for {address: bytes}. An application that reaches the walk
    region is sparse there - a .bin would have to carry every byte between."""
    lines = []
    for address, data in sorted(chunks.items()):
        for at in range(0, len(data), 16):
            row = data[at : at + 16]
            here = address + at
            record = bytes((len(row), here >> 8, here & 0xFF, 0)) + row
            lines.append(":" + (record + bytes((-sum(record) & 0xFF,))).hex())
    return "\n".join(lines + [":00000001FF"]) + "\n"


def main():
    device_bin, elf, mcu, hz, base_hex, page, baud, app_bin, tool, workdir = sys.argv[1:]
    page, baud = 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")
    # The ssd1306 board's own fuses, and the same boot section with reset left
    # on the application - the one difference that decides the refusal.
    bricking, harmless = "ffffffdc", "ffffffdd"

    def session(fuses):
        return pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, dump, fuses=fuses)

    def connected(device):
        port = pb.Port(device.pty, baud)
        loader = pb.Loader(port)
        loader.connect(25)
        return port, loader

    device = session(bricking)
    try:
        port, loader = connected(device)
        info = loader.info
        bootrst, bls_start = pb.mega_boot(info, bytes.fromhex(bricking))
        if not bootrst or bls_start >= info.base:
            fail(f"this profile cannot brick: BOOTRST {bootrst}, section at {bls_start:#06x}")
        port.close()

        # An application whose tail lands in the walk region, which is what the
        # bricked board's dump showed: live bytes at the section base.
        overgrown = os.path.join(workdir, "overgrown.hex")
        with open(overgrown, "w") as image:
            image.write(intel_hex({0: open(app_bin, "rb").read(), bls_start: bytes(page)}))

        out = pbsim.run_tool(tool, device.pty, baud, "--fuses", "--stay")
        if f"0x{bricking[6:]}" not in out:
            fail(f"the device did not answer the fuses it was given ({bricking}):\n{out}")

        try:
            pbsim.run_tool(tool, device.pty, baud, "--flash", overgrown, "--stay")
        except RuntimeError as refused:
            if "walk region" not in str(refused):
                fail(f"--flash failed, but not on the walk region: {refused}")
        else:
            fail("a plain --flash wrote an application into the reset walk region")

        port, loader = connected(device)
        if loader.read_flash(bls_start, page) != bytes((0xFF,)) * page:
            fail(f"the refused image was written to {bls_start:#06x} anyway")
        port.close()

        # The override is the escape hatch, not the default.
        pbsim.run_tool(tool, device.pty, baud, "--flash", overgrown, "--erase-flash", "--force", "--stay")
        port, loader = connected(device)
        if loader.read_flash(bls_start, page) != bytes(page):
            fail(f"--force did not write {bls_start:#06x}")
        port.close()
    finally:
        device.stop()

    # An image that stays below the region flashes with no override at all,
    # and so does the same overgrown image once reset boots the application.
    for fuses, image in ((bricking, app_bin), (harmless, overgrown)):
        device = session(fuses)
        try:
            pbsim.run_tool(tool, device.pty, baud, "--flash", image, "--erase-flash", "--stay")
        finally:
            device.stop()

    print("pbwalk: the walk region is refused by default, overridden by --force, and left alone otherwise")


main()