bootloader/test/tsbtest.py

#!/usr/bin/env python3
"""End-to-end TSB protocol test: spawn the simavr device, speak the TinySafeBoot
wire protocol over its pty (as the real host tools do), and actually flash it.

Usage: tsbtest.py <device_binary> <tsb.elf> <boot_base_hex>
Exits 0 if every scenario passes.
"""
import os
import subprocess
import sys
import time

import serial

CONFIRM = 0x21  # '!'
REQUEST = 0x3F  # '?'
KNOCK = 0x40    # '@'
PAGE = 128      # ATmega328P: 64 words


class Device:
    """The simavr runner, exposing UART0 as a pty. `config` seeds the config
    page (via the device's TSB_CONFIG hook) so the password gate and emergency
    erase are exercisable."""

    def __init__(self, binary, elf, boot_base, dump="/tmp/tsb_dump.bin", config=None):
        env = dict(os.environ)
        if config is not None:
            env["TSB_CONFIG"] = config
        self.proc = subprocess.Popen(
            [binary, elf, boot_base, dump],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env)
        self.dump = dump
        self.pty = None
        deadline = time.time() + 5
        while time.time() < deadline:
            line = self.proc.stdout.readline()
            if not line:
                break
            if line.startswith("TSB_PTY"):
                self.pty = line.split()[1]
                break
        if not self.pty:
            self.stop()
            raise RuntimeError("device did not report a pty")

    def stop(self):
        self.proc.terminate()
        try:
            self.proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.proc.kill()


class Host:
    """A faithful TSB host, per the wire protocol."""

    def __init__(self, pty):
        self.s = serial.Serial(pty, 115200, timeout=1.5)
        self.info = None

    def _read(self, n):
        data = self.s.read(n)
        if len(data) != n:
            raise AssertionError(f"expected {n} bytes, got {len(data)}: {data.hex()}")
        return data

    def activate(self):
        self.s.reset_input_buffer()
        self.s.write(b"@@@")
        reply = self._read(17)
        if reply[16] != CONFIRM:
            raise AssertionError(f"activation reply not '!'-terminated: {reply.hex()}")
        self.info = reply[:16]
        return self.info

    # Parsed info-block fields (host math from the spec).
    @property
    def pagesize(self):
        return self.info[9] * 2

    @property
    def appflash(self):
        return (self.info[10] | (self.info[11] << 8)) * 2

    @property
    def eeprom_size(self):
        return (self.info[12] | (self.info[13] << 8)) + 1

    def _expect(self, byte, what):
        r = self._read(1)
        if r[0] != byte:
            raise AssertionError(f"{what}: expected {byte:#x}, got {r.hex()}")

    # Host-paced page read ('f'/'e'): send '!', take a page, repeat; stop with
    # anything else, then the Mainloop '!'.
    def _read_pages(self, cmd, npages):
        self.s.write(cmd.encode())
        data = b""
        for _ in range(npages):
            self.s.write(bytes([CONFIRM]))
            data += self._read(PAGE)
        self.s.write(bytes([REQUEST]))  # stop
        self._expect(CONFIRM, f"{cmd} end")
        return data

    # Device-paced page write ('F'/'E'): device offers '?', host sends '!'+page,
    # or anything else to stop.
    def _write_pages(self, cmd, data):
        if len(data) % PAGE:
            data += b"\xff" * (PAGE - len(data) % PAGE)
        self.s.write(cmd.encode())
        for off in range(0, len(data), PAGE):
            self._expect(REQUEST, f"{cmd} '?'")
            self.s.write(bytes([CONFIRM]) + data[off:off + PAGE])
        self._expect(REQUEST, f"{cmd} trailing '?'")
        self.s.write(bytes([REQUEST]))  # stop
        self._expect(CONFIRM, f"{cmd} end")

    def write_flash(self, data):
        self._write_pages("F", data)

    def read_flash(self, npages):
        return self._read_pages("f", npages)

    def write_eeprom(self, data):
        self._write_pages("E", data)

    def read_eeprom(self, npages):
        return self._read_pages("e", npages)

    def read_config(self):
        self.s.write(b"c")
        page = self._read(PAGE)
        self._expect(CONFIRM, "c end")
        return page

    def write_config(self, data):
        assert len(data) == PAGE
        self.s.write(b"C")
        self._expect(REQUEST, "C '?'")
        self.s.write(bytes([CONFIRM]) + data)
        echo = self._read(PAGE)  # device echoes what it programmed
        self._expect(CONFIRM, "C end")
        return echo

    # Activation when the config page carries a password: 3x'@' then the
    # password bytes, then the info block + mainloop '!'.
    def activate_password(self, password):
        self.s.reset_input_buffer()
        self.s.write(bytes([KNOCK, KNOCK, KNOCK]) + password)
        reply = self._read(17)
        if reply[16] != CONFIRM:
            raise AssertionError(f"password activation not '!'-terminated: {reply.hex()}")
        self.info = reply[:16]
        return self.info

    # A 0 byte where a password byte is expected requests emergency erase; the
    # device asks for two confirmations, then wipes and returns to the mainloop.
    def emergency_erase(self):
        self.s.reset_input_buffer()
        self.s.write(bytes([KNOCK, KNOCK, KNOCK, 0x00]))
        self._expect(REQUEST, "emergency confirm 1")
        self.s.write(bytes([CONFIRM]))
        self._expect(REQUEST, "emergency confirm 2")
        self.s.write(bytes([CONFIRM]))
        self._expect(CONFIRM, "emergency mainloop ready")


def check(cond, msg):
    if not cond:
        raise AssertionError(msg)
    print(f"  ok: {msg}")


# A config page carrying a password "PW": appjump 0, timeout 0x40, password
# 0x50 0x57 terminated by 0xff.
PW_CONFIG = "0000405057ff"
PW_BYTES = bytes([0x50, 0x57])


def scenario_roundtrip(host):
    """Activation + info block + flash/EEPROM/config read-write round-trips, on
    a device with a blank (erased) config page - the usual no-password case."""
    info = host.activate()
    check(info[0:3] == b"TSB", f"magic 'TSB' (got {info[0:3]!r})")
    check(info[6:9] == bytes([0x1E, 0x95, 0x0F]), f"signature 1E 95 0F (got {info[6:9].hex()})")
    check(info[14] == info[15], f"device-type bytes 14==15 (got {info[14]:#x},{info[15]:#x})")
    check(host.pagesize == PAGE, f"page size {PAGE} (got {host.pagesize})")
    check(host.eeprom_size == 1024, f"eeprom size 1024 (got {host.eeprom_size})")
    print(f"  info: {info.hex()}  appflash={host.appflash} eeprom={host.eeprom_size}")

    app = bytes(range(256))  # two pages of known data
    host.write_flash(app)
    check(host.read_flash(2) == app, "flash round-trip 2 pages")

    edata = bytes((i * 7) & 0xFF for i in range(PAGE))
    host.write_eeprom(edata)
    check(host.read_eeprom(1) == edata, "eeprom round-trip 1 page")

    cfg = bytes([0x00, 0x00, 0x40]) + b"\xff" * (PAGE - 3)  # timeout 0x40, no password
    check(host.write_config(cfg) == cfg, "config write echoes the programmed page")
    check(host.read_config() == cfg, "config read-back matches")


def scenario_password(host):
    """A device whose config page carries a password activates only when the
    host sends it after the knock."""
    info = host.activate_password(PW_BYTES)
    check(info[0:3] == b"TSB", f"password activation returns the info block (got {info[0:3]!r})")


def scenario_emergency(host):
    """Emergency erase (password 0-byte + two confirms) wipes flash, EEPROM and
    the config page; the device stays alive in its boot section."""
    host.emergency_erase()
    check(host.read_config() == b"\xff" * PAGE, "config page wiped")
    check(host.read_flash(1) == b"\xff" * PAGE, "application flash wiped")
    check(host.read_eeprom(1) == b"\xff" * PAGE, "EEPROM wiped")


def main():
    binary, elf, boot_base = sys.argv[1], sys.argv[2], sys.argv[3]
    failures = []

    # Each group runs on its own freshly-reset device (simavr reloads the ELF,
    # so nothing persists between them); the password groups seed a config page.
    groups = [
        ("round-trip", None, scenario_roundtrip),
        ("password activation", PW_CONFIG, scenario_password),
        ("emergency erase", PW_CONFIG, scenario_emergency),
    ]
    for name, config, fn in groups:
        print(f"--- {name} ---")
        dev = Device(binary, elf, boot_base, config=config)
        try:
            fn(Host(dev.pty))
        except AssertionError as e:
            failures.append(f"{name}: {e}")
            print(f"  FAIL: {e}")
        finally:
            dev.stop()

    if failures:
        print(f"FAILED ({len(failures)})")
        return 1
    print("ALL PASS")
    return 0


if __name__ == "__main__":
    sys.exit(main())