tools: ota_update: improve edge case handling

src/infuse_iot/tools/gateway.py

@@ -75,7 +75,9 @@ def notification_broadcast(self, notification: ClientNotification):
         if self.server:
             self.server.broadcast(notification)
 
-    def query_device_key(self, infuse_id: int, cb_event: threading.Event | None = None):
+    def query_device_key(self, infuse_id: int, cb_event: threading.Event | None = None) -> bool:
+        """Query device key, returns True on success, False on failure"""
+
         def security_state_done(pkt: PacketReceived, _rc: int, response: bytes, challenge):
             decoded = defs.security_state.response.vla_from_buffer_copy(response)
             self.ddb.observe_security_state(
@@ -98,28 +100,32 @@ def public_keys_done(pkt: PacketReceived, rc: int, response: bytes, _):
                 if key.id == defs.rpc_enum_key_id.SECONDARY_REMOTE_PUBLIC_KEY:
                     self.ddb.observe_secondary_remote_public_key(infuse_id, bytes(key.key))
 
-        def run_cmd_pkt(cmd_pkt: PacketOutputRouted):
+        def run_cmd_pkt(cmd_pkt: PacketOutputRouted) -> bool:
             encrypted = cmd_pkt.to_serial(self.ddb)
             # Write to serial port
             Console.log_tx(cmd_pkt.ptype, len(encrypted))
             self.port.write(encrypted)
             if cb_event is not None:
                 # Wait for the response
-                cb_event.wait(1.0)
+                return cb_event.wait(1.0)
+            return True
 
         # Run security_state RPC
         challenge = random.randbytes(16)
         cmd_pkt = self.rpc.generate_addressed(
             infuse_id, defs.security_state.COMMAND_ID, challenge, Auth.NETWORK, security_state_done, challenge
         )
-        run_cmd_pkt(cmd_pkt)
+        if not run_cmd_pkt(cmd_pkt):
+            return False
 
         if self.ddb.has_local_root:
             # Query other public keys from the device
             cmd_pkt = self.rpc.generate_addressed(
                 infuse_id, defs.security_public_keys.COMMAND_ID, b"\x00", Auth.NETWORK, public_keys_done, None
             )
-            run_cmd_pkt(cmd_pkt)
+            if not run_cmd_pkt(cmd_pkt):
+                return False
+        return True
 
 
 class SerialRxThread(SignaledThread):
@@ -271,7 +277,9 @@ def _handle_epacket_send(self, req: GatewayRequestEpacketSend):
         for hop in routed.route:
             if hop.auth == Auth.DEVICE and not self._common.ddb.has_shared_key(hop.infuse_id):
                 cb_event = threading.Event()
-                self._common.query_device_key(hop.infuse_id, cb_event)
+                if not self._common.query_device_key(hop.infuse_id, cb_event):
+                    Console.log_error("Failed to query device key for {infuse_id:016x}")
+                    return
 
         # Encode and encrypt payload
         encrypted = routed.to_serial(self._common.ddb)
@@ -323,6 +331,11 @@ def _bt_connect_cb(self, _pkt: PacketReceived, rc: int, response: bytes, _):
             self._common.port.write(encrypted)
             return
 
+        if not self._common.ddb.has_shared_key(infuse_id):
+            # Pro-actively query key information
+            cb_event = threading.Event()
+            self._common.query_device_key(infuse_id, cb_event)
+
         # Notify connection success
         self._connected_notification(infuse_id)
 
@@ -341,7 +354,7 @@ def _handle_conn_request(self, req: GatewayRequestConnectionRequest):
 
         subs = 0
         bt_char = defs.rpc_enum_infuse_bt_characteristic
-        if req.data_types & req.DataType.COMMAND:
+        if req.data_types & req.DataType.COMMAND or not self._common.ddb.has_shared_key(req.infuse_id):
             subs |= bt_char.COMMAND
         if req.data_types & req.DataType.DATA:
             subs |= bt_char.DATA

src/infuse_iot/tools/ota_upgrade.py

@@ -324,12 +324,12 @@ def run(self):
                         source.infuse_id, GatewayRequestConnectionRequest.DataType.COMMAND, self._conn_timeout
                     ) as mtu:
                         # Set time on the remote device to keep keys in sync
-                        if not self.set_device_time(mtu, source.infuse_id):
-                            break
-                        if self._single_diff:
-                            self.run_file_copy(live, mtu, source)
-                        else:
-                            self.run_file_upload(live, mtu, source)
+                        if self.set_device_time(mtu, source.infuse_id):
+                            # Upload the patch file
+                            if self._single_diff:
+                                self.run_file_copy(live, mtu, source)
+                            else:
+                                self.run_file_upload(live, mtu, source)
 
                 except ConnectionRefusedError:
                     self.state_update(live, "Scanning")