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]

The ESP32 power cycles 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 bootloader 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 socket
import sys
import time

DEFAULT_PORT = 2323
PAGE_SIZE = 128  # ATmega328P flash page in bytes
FLASH_SIZE = 30720  # 32 KB minus the 2 KB bootloader
SIGNATURE = b"\x1e\x95\x0f"  # ATmega328P

# 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

SYNC_TIMEOUT = 10.0  # seconds to wait for the bootloader after connecting
REPLY_TIMEOUT = 1.0


class UploadError(Exception):
    pass


def read_hex(path):
    """Parse an Intel HEX file into a bytearray (gaps filled with 0xFF)."""
    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) > FLASH_SIZE:
        raise UploadError(f"firmware is {len(image)} bytes, the Nano has room for {FLASH_SIZE}")
    # pad to whole pages
    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):
        self.sock.settimeout(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 power cycle)."""
        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 UploadError(
            f"no answer from the Nano bootloader within {SYNC_TIMEOUT:.0f} s. "
            "Is the Nano TX -> ESP32 GPIO 16 wire connected?")

    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 upload(hex_path, host, port):
    image = read_hex(hex_path)
    print(f"Firmware: {hex_path} ({len(image)} bytes)")

    print(f"Connecting to the ESP32 flash bridge at {host}:{port} ...")
    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)

    with sock:
        boot = Bootloader(sock)
        print("Waiting for the Nano bootloader (the ESP32 power cycles the display) ...")
        boot.sync()

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

        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])

    print("Done. The ESP32 restarts the Nano with the new firmware.")


def main():
    if len(sys.argv) != 3:
        print(__doc__)
        return 2
    hex_path, target = sys.argv[1], sys.argv[2]
    host, _, port = target.partition(":")
    try:
        upload(hex_path, host, int(port) if port else DEFAULT_PORT)
    except UploadError as e:
        sys.stdout.flush()
        print(f"ERROR: {e}", file=sys.stderr)
        return 1
    return 0


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