bootloader/tools/pbrig.py

#!/usr/bin/env python3
"""Hardware rig driver for pureboot: an ISP programmer beside a serial link.

The simulated suites (`test/pb*.py`) prove the protocol; this drives the same
loader on real silicon, where the things a cycle-exact simulator cannot model
live — an RC oscillator off its nominal, a reset edge that has to come from
somewhere, a serial bridge with its own idea of what a baud is.

Nothing here knows a port name, a part or a programmer. Every deployment fact
arrives from the command line or the environment, so the same script serves any
board: see `Deployment`. As a module it is the reset/flash/talk primitives that
`pbhw.py` builds its acceptance suite from; as a command it is the handful of
one-shot operations worth having on a rig — most importantly `backup`, which is
the only thing standing between a fuse experiment and an unrecoverable part.

Two rig facts are encoded here because they are not guessable and cost a
session each to learn:

* **An ISP access resets the part**, and it runs again the moment the programmer
  releases it. That is the only reset edge available when the serial adapter's
  DTR is not wired to reset — so a loader session begins with an ISP touch and
  knocks immediately after, which is what `Rig.pureboot()` does.
* **avrdude splits `-U memory:op:file:format` on colons**, so a Windows path's
  drive letter breaks the spec. Every file argument is therefore passed as a
  bare filename with avrdude run in that file's own directory.
"""

from __future__ import annotations

import argparse
import dataclasses
import importlib.util
import os
import pathlib
import subprocess
import sys
import time

HERE = pathlib.Path(__file__).resolve().parent
DEFAULT_PUREBOOT = HERE.parent / "pureboot" / "pureboot.py"

# Memories worth capturing before an experiment, and the format each is read in.
# Fuses and lock are per-part: a part without an extended fuse simply fails that
# one read, which `backup` reports and steps over rather than aborting on.
BACKUP_MEMORIES = (
    ("flash", "i", "hex"),
    ("flash", "r", "bin"),
    ("eeprom", "i", "hex"),
    ("eeprom", "r", "bin"),
    ("lfuse", "h", "hex"),
    ("hfuse", "h", "hex"),
    ("efuse", "h", "hex"),
    ("lock", "h", "hex"),
    ("calibration", "h", "hex"),
)


class Error(Exception):
    pass


def bitclock_for(hz: int) -> str:
    """A safe ISP bitclock for a part *currently running* at `hz`.

    SCK must stay under a quarter of the target clock, so the bitclock follows
    the clock in force — not the one about to be fused in. Halving that ceiling
    again costs nothing on a link that moves a few hundred bytes and buys margin
    against an oscillator that is already known to be off its nominal.
    """
    ceiling = hz // 8
    for candidate in (1000, 4000, 8000, 32000, 125000, 400000):
        if candidate <= ceiling:
            best = candidate
        else:
            break
    else:
        best = 400000
    if ceiling < 1000:
        raise Error(f"a part at {hz} Hz is too slow to reach over ISP safely")
    return f"{best // 1000}kHz"


@dataclasses.dataclass
class Deployment:
    """Everything about one board. No default names a real device."""

    port: str = ""              # serial device the loader speaks on
    baud: int = 57600           # host rate; for autobaud, the rate to drive
    autobaud: bool = False      # send the calibration pulse instead of p+b
    programmer: str = ""        # avrdude -c
    part: str = ""              # avrdude -p
    avrdude: str = "avrdude"
    bitclock: str = "125kHz"    # see bitclock_for()
    pureboot: pathlib.Path = DEFAULT_PUREBOOT
    wait: int = 12              # seconds the host keeps knocking

    @classmethod
    def from_env(cls) -> "Deployment":
        """Environment defaults, so a rig's facts live in one place per machine."""
        return cls(
            port=os.environ.get("PUREBOOT_PORT", ""),
            baud=int(os.environ.get("PUREBOOT_BAUD", "57600")),
            autobaud=os.environ.get("PUREBOOT_AUTOBAUD", "") not in ("", "0"),
            programmer=os.environ.get("PUREBOOT_PROGRAMMER", ""),
            part=os.environ.get("PUREBOOT_PART", ""),
            avrdude=os.environ.get("AVRDUDE", "avrdude"),
            bitclock=os.environ.get("PUREBOOT_BITCLOCK", "125kHz"),
            pureboot=pathlib.Path(os.environ.get("PUREBOOT_TOOL", str(DEFAULT_PUREBOOT))),
        )

    @staticmethod
    def add_arguments(parser: argparse.ArgumentParser) -> None:
        """Deployment flags, shared by this tool and pbhw.py."""
        env = Deployment.from_env()
        parser.add_argument("--port", default=env.port, help="serial device the loader speaks on")
        parser.add_argument("--baud", type=int, default=env.baud,
                            help="host rate (for autobaud, the rate to drive)")
        parser.add_argument("--autobaud", action="store_true", default=env.autobaud,
                            help="send the calibration pulse instead of the p+b knock")
        parser.add_argument("--programmer", default=env.programmer, help="avrdude -c, e.g. atmelice_isp")
        parser.add_argument("--part", default=env.part, help="avrdude -p, e.g. t13 or m328p")
        parser.add_argument("--avrdude", default=env.avrdude, help="path to avrdude")
        parser.add_argument("--bitclock", default=env.bitclock, help="ISP bitclock, e.g. 125kHz or 8kHz")
        parser.add_argument("--pureboot", type=pathlib.Path, default=env.pureboot,
                            help="path to pureboot.py")
        parser.add_argument("--wait", type=int, default=env.wait, help="seconds to keep knocking")

    @classmethod
    def from_args(cls, args: argparse.Namespace) -> "Deployment":
        return cls(port=args.port, baud=args.baud, autobaud=args.autobaud,
                   programmer=args.programmer, part=args.part, avrdude=args.avrdude,
                   bitclock=args.bitclock, pureboot=args.pureboot, wait=args.wait)


def load_pureboot(path: pathlib.Path = DEFAULT_PUREBOOT):
    """The host tool as a module — its Port and Loader, not a subprocess.

    Used where a subprocess cannot express what is needed: a poke followed by a
    peek in the *same* session, or a raw read at an arbitrary baud.
    """
    spec = importlib.util.spec_from_file_location("pureboot", path)
    if spec is None or spec.loader is None:
        raise Error(f"cannot load the host tool from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


class Rig:
    """One board: its programmer on one side, its serial link on the other."""

    def __init__(self, deployment: Deployment):
        self.d = deployment
        if not deployment.programmer or not deployment.part:
            raise Error("a rig needs --programmer and --part")

    # ------------------------------------------------------------- programmer

    def avrdude(self, *args: str, cwd: pathlib.Path | None = None,
                bitclock: str | None = None, timeout: int = 300) -> subprocess.CompletedProcess:
        command = [self.d.avrdude, "-c", self.d.programmer, "-p", self.d.part,
                   "-B", bitclock or self.d.bitclock, *args]
        return subprocess.run(command, capture_output=True, text=True,
                              cwd=None if cwd is None else str(cwd), timeout=timeout)

    @staticmethod
    def _ok(result: subprocess.CompletedProcess) -> bool:
        return result.returncode == 0

    def reset(self, bitclock: str | None = None) -> None:
        """An ISP access, which resets the part; it runs when avrdude exits."""
        self.avrdude("-U", "signature:r:-:h", bitclock=bitclock)

    def signature(self, bitclock: str | None = None) -> str:
        result = self.avrdude("-U", "signature:r:-:h", bitclock=bitclock)
        for line in reversed(result.stdout.splitlines()):
            if line.strip().startswith("0x"):
                return line.strip()
        raise Error(f"no signature read: {(result.stderr or result.stdout).strip()[:200]}")

    def read_memory(self, memory: str, destination: pathlib.Path, fmt: str = "r",
                    bitclock: str | None = None) -> bool:
        """Read `memory` into `destination`, whose directory avrdude runs in."""
        destination = pathlib.Path(destination).resolve()
        destination.parent.mkdir(parents=True, exist_ok=True)
        result = self.avrdude("-U", f"{memory}:r:{destination.name}:{fmt}",
                              cwd=destination.parent, bitclock=bitclock)
        # A memory the part does not have (a tiny's extended fuse) leaves avrdude
        # happy and the file empty. An empty capture is a miss, not a backup.
        return self._ok(result) and destination.exists() and destination.stat().st_size > 0

    def write_memory(self, memory: str, source: pathlib.Path, fmt: str = "i",
                     erase: bool = False, bitclock: str | None = None) -> bool:
        source = pathlib.Path(source).resolve()
        args = ["-U", f"{memory}:w:{source.name}:{fmt}"]
        if erase:
            args.insert(0, "-e")
        result = self.avrdude(*args, cwd=source.parent, bitclock=bitclock)
        return "verified" in (result.stdout + result.stderr)

    def flash_hex(self, image: pathlib.Path, erase: bool = True,
                  bitclock: str | None = None) -> bool:
        return self.write_memory("flash", image, "i", erase=erase, bitclock=bitclock)

    def read_fuses(self, bitclock: str | None = None) -> dict[str, str]:
        out: dict[str, str] = {}
        for fuse in ("lfuse", "hfuse", "efuse", "lock"):
            result = self.avrdude("-U", f"{fuse}:r:-:h", bitclock=bitclock)
            values = [l.strip() for l in result.stdout.splitlines() if l.strip().startswith("0x")]
            if values:
                out[fuse] = values[-1]
        return out

    def write_fuses(self, bitclock: str | None = None, **fuses: str) -> bool:
        """Write named fuses. A fuse change moves the clock the *next* access is
        timed against, so pass a bitclock safe for both sides of the change."""
        args: list[str] = []
        for name, value in fuses.items():
            args += ["-U", f"{name}:w:{value}:m"]
        if not args:
            return True
        result = self.avrdude(*args, bitclock=bitclock)
        text = result.stdout + result.stderr
        return "verified" in text or "written" in text

    # ------------------------------------------------------------ backup

    def backup(self, directory: pathlib.Path, prefix: str = "") -> dict[str, bool]:
        """Capture every memory worth keeping, then prove it by a second read.

        A backup nobody verified is a guess. Each memory is read twice and the
        two reads compared; a mismatch is reported rather than quietly stored.
        """
        directory = pathlib.Path(directory).resolve()
        directory.mkdir(parents=True, exist_ok=True)
        stem = prefix or self.d.part
        status: dict[str, bool] = {}
        for memory, fmt, extension in BACKUP_MEMORIES:
            name = f"{stem}-{memory}.{extension}"
            if not self.read_memory(memory, directory / name, fmt):
                status[f"{memory}.{extension}"] = False
                continue
            if extension == "bin":  # only the raw form is worth comparing byte-wise
                again = directory / f".{name}.again"
                self.read_memory(memory, again, fmt)
                same = again.exists() and again.read_bytes() == (directory / name).read_bytes()
                again.unlink(missing_ok=True)
                status[f"{memory}.{extension}"] = same
            else:
                status[f"{memory}.{extension}"] = True
        return status

    # ------------------------------------------------------------ serial link

    def pureboot(self, *args: str, reset_first: bool = True, baud: int | None = None,
                 autobaud: bool | None = None, timeout: int = 300,
                 bitclock: str | None = None) -> tuple[int, str]:
        """Reset, then knock immediately — see the module docstring.

        Returns the host tool's exit status and its combined output, so a caller
        can assert on what it printed as well as on whether it succeeded.
        """
        if reset_first:
            self.reset(bitclock=bitclock)
        command = [sys.executable, str(self.d.pureboot), "--port", self.d.port,
                   "--baud", str(self.d.baud if baud is None else baud),
                   "--wait", str(self.d.wait)]
        if self.d.autobaud if autobaud is None else autobaud:
            command.append("--autobaud")
        command += [str(a) for a in args]
        try:
            result = subprocess.run(command, capture_output=True, text=True, timeout=timeout)
        except subprocess.TimeoutExpired as expired:
            return 99, f"TIMEOUT after {timeout}s\n{expired.stdout or ''}{expired.stderr or ''}"
        return result.returncode, (result.stdout or "") + (result.stderr or "")

    def capture(self, seconds: float = 2.0, baud: int | None = None) -> bytes:
        """Listen to whatever the board is saying, at an arbitrary rate.

        Opening the port does not reset a board whose DTR is unwired, so this can
        sample a running application repeatedly without disturbing it — which is
        what makes the rate sweep below possible.
        """
        module = load_pureboot(self.d.pureboot)
        port = module.Port(self.d.port, self.d.baud if baud is None else baud)
        try:
            data = b""
            deadline = time.monotonic() + seconds
            while time.monotonic() < deadline:
                chunk = port.read_available(0.2)
                if chunk:
                    data += chunk
            return data
        finally:
            try:
                port.close()
            except Exception:
                pass


def measure_rate(rig: Rig, marker: bytes, built_baud: int, nominal_hz: int | None = None,
                 span_percent: float = 12.0, step_percent: float = 0.5,
                 seconds: float = 0.75) -> dict:
    """Find a transmitting board's true bit rate, using only the serial port.

    The board must be emitting something recognisable at a *fixed* cycles-per-bit
    — `test/pbapp.cpp` built with PUREBOOT_HEARTBEAT does. Since its bit timing is
    a cycle count, its wire rate scales with its actual clock, so the host rates
    at which `marker` still decodes bracket that rate; the centre of the band is
    the answer, and with the clock the image was built for it gives the real one.

    This is the measurement that turns "the loader is silent, so the wiring must
    be wrong" into a number, and it needs no instrument beyond the adapter
    already attached.
    """
    steps = int(span_percent / step_percent)
    clean: list[int] = []
    samples: list[tuple[int, int, bool]] = []
    for index in range(-steps, steps + 1):
        baud = int(round(built_baud * (1 + index * step_percent / 100.0)))
        if baud <= 0:
            continue
        data = rig.capture(seconds=seconds, baud=baud)
        hit = marker in data
        samples.append((baud, len(data), hit))
        if hit:
            clean.append(baud)
    result: dict = {"samples": samples, "clean": clean, "built_baud": built_baud}
    if clean:
        low, high = min(clean), max(clean)
        centre = (low + high) / 2.0
        result |= {"low": low, "high": high, "centre": centre,
                   "half_width_percent": (high - low) / 2.0 / centre * 100.0,
                   "error_percent": (centre / built_baud - 1.0) * 100.0}
        if nominal_hz:
            result["measured_hz"] = nominal_hz * centre / built_baud
    return result


# ------------------------------------------------------------------- command


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="pureboot hardware rig: ISP reset/flash beside the serial link")
    Deployment.add_arguments(parser)
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("signature", help="read the part signature over ISP")
    sub.add_parser("reset", help="reset the part (an ISP access) and let it run")
    sub.add_parser("fuses", help="read the fuse and lock bytes")

    p = sub.add_parser("flash", help="program a hex image over ISP")
    p.add_argument("image", type=pathlib.Path)
    p.add_argument("--no-erase", action="store_true", help="do not chip-erase first")

    p = sub.add_parser("backup", help="capture and verify every memory")
    p.add_argument("directory", type=pathlib.Path)
    p.add_argument("--prefix", default="", help="filename stem (default: the part name)")

    p = sub.add_parser("rate", help="measure the board's true bit rate and clock")
    p.add_argument("--marker", default="APP", help="text the board emits (default: APP)")
    p.add_argument("--built-baud", type=int, required=True,
                   help="the baud the running image was built for")
    p.add_argument("--nominal-hz", type=int, default=0,
                   help="the clock the image was built for, to report the real one")
    p.add_argument("--span", type=float, default=12.0, help="sweep +-this many percent")
    p.add_argument("--step", type=float, default=0.5, help="sweep step in percent")
    p.add_argument("--verbose", action="store_true", help="print every step")

    p = sub.add_parser("bitclock", help="a safe ISP bitclock for a clock in force")
    p.add_argument("hz", type=int)

    args = parser.parse_args(argv)

    if args.command == "bitclock":
        print(bitclock_for(args.hz))
        return 0

    rig = Rig(Deployment.from_args(args))

    if args.command == "signature":
        print(rig.signature())
    elif args.command == "reset":
        rig.reset()
        print("reset")
    elif args.command == "fuses":
        for name, value in rig.read_fuses().items():
            print(f"{name:<6} {value}")
    elif args.command == "flash":
        ok = rig.flash_hex(args.image, erase=not args.no_erase)
        print(f"{args.image.name}: {'verified' if ok else 'FAILED'}")
        return 0 if ok else 1
    elif args.command == "backup":
        status = rig.backup(args.directory, args.prefix)
        for name, ok in status.items():
            print(f"  {'ok  ' if ok else 'FAIL'} {name}")
        missing = [n for n, ok in status.items() if not ok]
        # Fuses a part does not have are expected misses, not failures.
        fatal = [n for n in missing if not n.startswith(("efuse", "calibration"))]
        print(f"\n{len(status) - len(missing)}/{len(status)} captured into {args.directory}")
        return 1 if fatal else 0
    elif args.command == "rate":
        result = measure_rate(rig, args.marker.encode(), args.built_baud,
                              args.nominal_hz or None, args.span, args.step)
        if args.verbose:
            for baud, size, hit in result["samples"]:
                print(f"  {baud:7d} Bd  {size:5d} B  {'MARKER' if hit else ''}")
        if not result["clean"]:
            print(f"no capture contained {args.marker!r} at any rate — is the board "
                  f"transmitting, and on the pin this port is wired to?")
            return 1
        print(f"clean band  {result['low']}..{result['high']} Bd")
        print(f"centre      {result['centre']:.0f} Bd "
              f"(+-{result['half_width_percent']:.1f} %)")
        print(f"vs built    {result['built_baud']} Bd ({result['error_percent']:+.1f} %)")
        if "measured_hz" in result:
            print(f"true clock  {result['measured_hz'] / 1e6:.3f} MHz")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Error as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)