#!/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 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: 3×'@' 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())