feat: Add ESP32 WiFi Unified OTA update support

meshtastic/__main__.py

@@ -37,6 +37,7 @@ try:
 except ImportError as e:
     have_test = False
 
+import meshtastic.ota
 import meshtastic.util
 import meshtastic.serial_interface
 import meshtastic.tcp_interface
@@ -60,7 +61,7 @@ except ImportError as e:
     have_powermon = False
     powermon_exception = e
     meter = None
-from meshtastic.protobuf import channel_pb2, config_pb2, portnums_pb2, mesh_pb2
+from meshtastic.protobuf import admin_pb2, channel_pb2, config_pb2, portnums_pb2, mesh_pb2
 from meshtastic.version import get_active_version
 
 logger = logging.getLogger(__name__)
@@ -452,6 +453,41 @@ def onConnected(interface):
             waitForAckNak = True
             interface.getNode(args.dest, False, **getNode_kwargs).rebootOTA()
 
+        if args.ota_update:
+            closeNow = True
+            waitForAckNak = True
+
+            if not isinstance(interface, meshtastic.tcp_interface.TCPInterface):
+                meshtastic.util.our_exit(
+                    "Error: OTA update currently requires a TCP connection to the node (use --host)."
+                )
+
+            ota = meshtastic.ota.ESP32WiFiOTA(args.ota_update, interface.hostname)
+
+            print(f"Triggering OTA update on {interface.hostname}...")
+            interface.getNode(args.dest, False, **getNode_kwargs).startOTA(
+                mode=admin_pb2.OTA_WIFI,
+                hash=ota.hash_bytes()
+            )
+
+            print("Waiting for device to reboot into OTA mode...")
+            time.sleep(5)
+
+            retries = 5
+            while retries > 0:
+                try:
+                    ota.update()
+                    break
+
+                except Exception as e:
+                    retries -= 1
+                    if retries == 0:
+                        meshtastic.util.our_exit(f"\nOTA update failed: {e}")
+
+                    time.sleep(2)
+
+            print("\nOTA update completed successfully!")
+                
         if args.enter_dfu:
             closeNow = True
             waitForAckNak = True
@@ -1904,10 +1940,17 @@ def addRemoteAdminArgs(parser: argparse.ArgumentParser) -> argparse.ArgumentPars
 
     group.add_argument(
         "--reboot-ota",
-        help="Tell the destination node to reboot into factory firmware (ESP32)",
+        help="Tell the destination node to reboot into factory firmware (ESP32, firmware version <2.7.18)",
         action="store_true",
     )
 
+    group.add_argument(
+        "--ota-update",
+        help="Perform a OTA update on the destination node (ESP32, firmware version >=2.7.18, WiFi/TCP only for now). Specify the path to the firmware file.",
+        metavar="FIRMWARE_FILE",
+        action="store",
+    )
+
     group.add_argument(
         "--enter-dfu",
         help="Tell the destination node to enter DFU mode (NRF52)",

meshtastic/node.py

@@ -654,7 +654,7 @@ class Node:
         return self._sendAdmin(p, onResponse=onResponse)
 
     def rebootOTA(self, secs: int = 10):
-        """Tell the node to reboot into factory firmware."""
+        """Tell the node to reboot into factory firmware (firmware < 2.7.18)."""
         self.ensureSessionKey()
         p = admin_pb2.AdminMessage()
         p.reboot_ota_seconds = secs
@@ -667,6 +667,22 @@ class Node:
             onResponse = self.onAckNak
         return self._sendAdmin(p, onResponse=onResponse)
 
+    def startOTA(
+        self,
+        mode: admin_pb2.OTAMode.ValueType,
+        hash: bytes,
+    ):
+        """Tell the node to start OTA mode (firmware >= 2.7.18)."""
+        if self != self.iface.localNode:
+            raise Exception("startOTA only possible in local node")
+
+        self.ensureSessionKey()
+        p = admin_pb2.AdminMessage()
+        p.ota_request.reboot_ota_mode=mode
+        p.ota_request.ota_hash=hash
+
+        return self._sendAdmin(p)
+
     def enterDFUMode(self):
         """Tell the node to enter DFU mode (NRF52)."""
         self.ensureSessionKey()

meshtastic/ota.py

@@ -0,0 +1,122 @@
+import os 
+import hashlib
+import socket
+import logging
+from typing import Optional, Callable
+
+
+logger = logging.getLogger(__name__)
+
+
+def _file_sha256(filename: str):
+    """Calculate SHA256 hash of a file."""
+    sha256_hash = hashlib.sha256()
+    
+    with open(filename, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+
+    return sha256_hash
+
+class ESP32WiFiOTA:
+    """ESP32 WiFi Unified OTA updates."""
+
+    def __init__(self, filename: str, hostname: str, port: int = 3232):
+        self._filename = filename
+        self._hostname = hostname
+        self._port = port
+        self._socket: Optional[socket.socket] = None
+
+        if not os.path.exists(self._filename):
+            raise Exception(f"File {self._filename} does not exist")
+
+        self._file_hash = _file_sha256(self._filename)
+
+    def _read_line(self) -> str:
+        """Read a line from the socket."""
+        if not self._socket:
+            raise Exception("Socket not connected")
+
+        line = b""
+        while not line.endswith(b"\n"):
+            char = self._socket.recv(1)
+            
+            if not char:
+                raise Exception("Connection closed while waiting for response")
+
+            line += char
+
+        return line.decode("utf-8").strip()
+
+    def hash_bytes(self) -> bytes:
+        """Return the hash as bytes."""
+        return self._file_hash.digest()
+
+    def hash_hex(self) -> str:
+        """Return the hash as a hex string."""
+        return self._file_hash.hexdigest()
+
+    def update(self, progress_callback: Optional[Callable[[int, int], None]] = None):
+        """Perform the OTA update."""
+        with open(self._filename, "rb") as f:
+            data = f.read()
+        size = len(data)
+
+        logger.info(f"Starting OTA update with {self._filename} ({size} bytes, hash {self.hash_hex()})")
+
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._socket.settimeout(15)
+        try:
+            self._socket.connect((self._hostname, self._port))
+            logger.debug(f"Connected to {self._hostname}:{self._port}")
+
+            # Send start command
+            self._socket.sendall(f"OTA {size} {self.hash_hex()}\n".encode("utf-8"))
+
+            # Wait for OK from the device
+            while True:
+                response = self._read_line()
+                if response == "OK":
+                    break
+                elif response == "ERASING":
+                    logger.info("Device is erasing flash...")
+                elif response.startswith("ERR "):
+                    raise Exception(f"Device reported error: {response}")
+                else:
+                    logger.warning(f"Unexpected response: {response}")
+
+            # Stream firmware
+            sent_bytes = 0
+            chunk_size = 1024
+            while sent_bytes < size:
+                chunk = data[sent_bytes : sent_bytes + chunk_size]
+                self._socket.sendall(chunk)
+                sent_bytes += len(chunk)
+
+                if progress_callback:
+                    progress_callback(sent_bytes, size)
+                else:
+                    print(f"[{sent_bytes / size * 100:5.1f}%] Sent {sent_bytes} of {size} bytes...", end="\r")
+
+            if not progress_callback:
+                print()
+
+            # Wait for OK from device
+            logger.info("Firmware sent, waiting for verification...")
+            while True:
+                response = self._read_line()
+
+                if response == "OK":
+                    logger.info("OTA update completed successfully!")
+                    break
+                elif response == "ACK":
+                    continue
+                elif response.startswith("ERR "):
+                    raise Exception(f"OTA update failed: {response}")
+                else:
+                    logger.warning(f"Unexpected final response: {response}")
+
+        finally:
+            if self._socket:
+                self._socket.close()
+                self._socket = None