Flipdot-display/tools/nano_net_upload.py

#!/usr/bin/env python3
"""Flash the Arduino Nano over WiFi, through the ESP32's flash bridge.

Usage:  python nano_net_upload.py firmware.hex HOST[:PORT] [--baud auto|57600|115200]

--baud is the Nano bootloader's baud rate:
  57600   old Nano bootloader (ATmegaBOOT)
  115200  Optiboot ("new bootloader"); needs the ESP32 reset wire
  auto    try 57600 first, then 115200 (default)

The ESP32 restarts the Nano when we connect (so its bootloader starts) and then passes
bytes between this TCP connection and the Nano's serial port. This script talks the
STK500v1 protocol of the Arduino bootloaders directly. avrdude can't be used here: its
network ports are unreliable on Windows.

Normally run by PlatformIO: pio run -e nano_ota -t upload
"""

import argparse
import socket
import sys
import time

DEFAULT_PORT = 2323
PAGE_SIZE = 128  # ATmega328P flash page in bytes
SIGNATURE = b"\x1e\x95\x0f"  # ATmega328P

# Bootloaders: baud rate -> (name, flash space left for the firmware)
BOOTLOADERS = {
    57600: ("old bootloader", 32768 - 2048),
    115200: ("Optiboot", 32768 - 512),
}
AUTO_ORDER = [57600, 115200]

# STK500v1
STK_OK = 0x10
STK_INSYNC = 0x14
CRC_EOP = 0x20
GET_SYNC = 0x30
ENTER_PROGMODE = 0x50
LEAVE_PROGMODE = 0x51
LOAD_ADDRESS = 0x55
PROG_PAGE = 0x64
READ_PAGE = 0x74
READ_SIGN = 0x75

BRIDGE_TIMEOUT = 5.0  # wait for the ESP32's OK (it power cycles the Nano first)
SYNC_TIMEOUT = 4.0  # wait for the bootloader to answer after the restart
REPLY_TIMEOUT = 1.0


class UploadError(Exception):
    pass


class NoBootloader(UploadError):
    """The bootloader didn't answer (wrong baud rate, missing wire, ...)."""


def read_hex(path):
    """Parse an Intel HEX file into a bytearray (gaps filled with 0xFF), padded to whole pages."""
    image = bytearray()
    base = 0
    with open(path) as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            if not line.startswith(":"):
                raise UploadError(f"{path}:{line_no}: not an Intel HEX line")
            raw = bytes.fromhex(line[1:])
            if sum(raw) & 0xFF:
                raise UploadError(f"{path}:{line_no}: checksum error")
            count, addr, rtype, data = raw[0], (raw[1] << 8) | raw[2], raw[3], raw[4:4 + raw[0]]
            if len(data) != count:
                raise UploadError(f"{path}:{line_no}: length mismatch")
            if rtype == 0x00:  # data
                start = base + addr
                if start + count > len(image):
                    image.extend(b"\xff" * (start + count - len(image)))
                image[start:start + count] = data
            elif rtype == 0x01:  # end of file
                break
            elif rtype == 0x02:  # extended segment address
                base = ((data[0] << 8) | data[1]) << 4
            elif rtype == 0x04:  # extended linear address
                base = ((data[0] << 8) | data[1]) << 16
    if len(image) % PAGE_SIZE:
        image.extend(b"\xff" * (PAGE_SIZE - len(image) % PAGE_SIZE))
    return image


class Bootloader:
    def __init__(self, sock):
        self.sock = sock

    def _recv(self, n, timeout=REPLY_TIMEOUT):
        data = b""
        deadline = time.monotonic() + timeout
        while len(data) < n:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            self.sock.settimeout(remaining)
            try:
                chunk = self.sock.recv(n - len(data))
            except socket.timeout:
                break
            if not chunk:
                raise UploadError("connection closed by the ESP32")
            data += chunk
        return data

    def _drain(self):
        """Throw away everything that is still arriving (e.g. answers to earlier syncs)."""
        while self._recv(256, timeout=0.1):
            pass

    def command(self, payload, reply_len=0):
        """Send a command and return its reply data (between INSYNC and OK)."""
        self.sock.sendall(bytes(payload) + bytes([CRC_EOP]))
        reply = self._recv(reply_len + 2)
        if len(reply) != reply_len + 2 or reply[0] != STK_INSYNC or reply[-1] != STK_OK:
            raise UploadError(f"bad reply to command 0x{payload[0]:02x}: {reply.hex() or 'nothing'}")
        return reply[1:-1]

    def sync(self):
        """Keep sending GET_SYNC until the bootloader answers (it starts after the restart)."""
        deadline = time.monotonic() + SYNC_TIMEOUT
        while time.monotonic() < deadline:
            self.sock.sendall(bytes([GET_SYNC, CRC_EOP]))
            if self._recv(2, timeout=0.2) == bytes([STK_INSYNC, STK_OK]):
                # Answers to earlier syncs may still be on their way: clear them, then
                # sync once more so the replies line up with our commands.
                self._drain()
                self.command([GET_SYNC])
                return
        raise NoBootloader(f"no answer from the Nano bootloader within {SYNC_TIMEOUT:.0f} s")

    def load_address(self, byte_addr):
        word = byte_addr // 2
        self.command([LOAD_ADDRESS, word & 0xFF, word >> 8])

    def write_page(self, byte_addr, data):
        self.load_address(byte_addr)
        self.command([PROG_PAGE, len(data) >> 8, len(data) & 0xFF, ord("F")] + list(data))

    def read_page(self, byte_addr, length):
        self.load_address(byte_addr)
        return self.command([READ_PAGE, length >> 8, length & 0xFF, ord("F")], reply_len=length)


def progress(label, done, total, step=PAGE_SIZE):
    """Print a line every 10% (PlatformIO shows each update on its own line)."""
    percent = 100 * done // total
    previous = 100 * (done - step) // total
    if done == step or done == total or percent // 10 != previous // 10:
        print(f"{label} {percent:3d}%  ({done}/{total} bytes)", flush=True)


def connect(host, port, baud):
    """Connect to the ESP32 bridge and have it restart the Nano into its bootloader."""
    try:
        sock = socket.create_connection((host, port), timeout=10)
    except OSError as e:
        raise UploadError(f"can't connect to {host}:{port}: {e}")
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

    sock.sendall(f"FLASH {baud}\n".encode())
    reply = b""
    deadline = time.monotonic() + BRIDGE_TIMEOUT
    while not reply.endswith(b"\n") and time.monotonic() < deadline:
        sock.settimeout(max(deadline - time.monotonic(), 0.01))
        try:
            chunk = sock.recv(64)
        except socket.timeout:
            break
        if not chunk:
            break
        reply += chunk
    if reply != b"OK\n":
        sock.close()
        raise UploadError(
            f"the ESP32 flash bridge didn't accept the request (reply: {reply!r}). "
            "Is the ESP32 running up-to-date firmware?")
    return sock


def flash(sock, image, baud):
    name, flash_size = BOOTLOADERS[baud]
    if len(image) > flash_size:
        raise UploadError(f"firmware is {len(image)} bytes, but with the {name} the Nano has room for {flash_size}")

    boot = Bootloader(sock)
    signature = boot.command([READ_SIGN], reply_len=3)
    if signature != SIGNATURE:
        raise UploadError(f"unexpected device signature {signature.hex()} (expected ATmega328P)")
    print(f"Found ATmega328P, bootloader: {name} ({baud} baud)")

    boot.command([ENTER_PROGMODE])

    for addr in range(0, len(image), PAGE_SIZE):
        boot.write_page(addr, image[addr:addr + PAGE_SIZE])
        progress("Writing  ", addr + PAGE_SIZE, len(image))

    for addr in range(0, len(image), PAGE_SIZE):
        expected = image[addr:addr + PAGE_SIZE]
        actual = boot.read_page(addr, PAGE_SIZE)
        if actual != expected:
            for i, (a, b) in enumerate(zip(actual, expected)):
                if a != b:
                    raise UploadError(f"verify failed at 0x{addr + i:04x}: read 0x{a:02x}, expected 0x{b:02x}")
        progress("Verifying", addr + PAGE_SIZE, len(image))

    boot.command([LEAVE_PROGMODE])


def upload(hex_path, host, port, bauds):
    image = read_hex(hex_path)
    print(f"Firmware: {hex_path} ({len(image)} bytes)")

    for i, baud in enumerate(bauds):
        print(f"Connecting to the ESP32 flash bridge at {host}:{port}, "
              f"restarting the Nano into its bootloader at {baud} baud ...", flush=True)
        with connect(host, port, baud) as sock:
            try:
                Bootloader(sock).sync()
            except NoBootloader:
                if i + 1 < len(bauds):
                    print(f"No answer at {baud} baud, trying {bauds[i + 1]} ...", flush=True)
                    continue
                raise NoBootloader(
                    f"no answer from the Nano bootloader at {', '.join(map(str, bauds))} baud. "
                    "Check the Nano TX -> ESP32 GPIO 16 wire. A Nano with Optiboot also needs "
                    "the reset wire (nano_reset in src/esp32/libs.h).")
            flash(sock, image, baud)
            print("Done. The ESP32 restarts the Nano with the new firmware.")
            return


def main():
    parser = argparse.ArgumentParser(description="Flash the Nano over WiFi through the ESP32.")
    parser.add_argument("hex", help="firmware .hex file")
    parser.add_argument("target", help=f"ESP32 address, HOST[:PORT] (default port {DEFAULT_PORT})")
    parser.add_argument("--baud", default="auto", choices=["auto"] + [str(b) for b in BOOTLOADERS],
                        help="bootloader baud rate (default: auto)")
    args = parser.parse_args()

    host, _, port = args.target.partition(":")
    bauds = AUTO_ORDER if args.baud == "auto" else [int(args.baud)]
    try:
        upload(args.hex, host, int(port) if port else DEFAULT_PORT, bauds)
    except UploadError as e:
        sys.stdout.flush()
        print(f"ERROR: {e}", file=sys.stderr)
        return 1
    return 0


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