Enhanced firmware OTA update (#98)

* Added firmware update exposed to user

* Added firmware update error attribute

* Removed random dummy failed on update checksum verification

* Updated firmware update example

Author: LeoGranierWiifor

examples/device/firmware_update.py

@@ -14,10 +14,26 @@
 
 import time
 import logging
-from tb_device_mqtt import TBDeviceMqttClient, FW_STATE_ATTR
+from random import randint
+from tb_device_mqtt import TBDeviceMqttClient, TBFirmwareState, FW_STATE_ATTR, FW_TITLE_ATTR
 
 logging.basicConfig(level=logging.INFO)
 
+def on_firmware_received_example(client, firmware_data, version_to):
+    client.update_firmware_info(state = TBFirmwareState.UPDATING)
+    time.sleep(1)
+
+    with open(client.firmware_info.get(FW_TITLE_ATTR), "wb") as firmware_file:
+        firmware_file.write(firmware_data)
+
+    random_value = randint(0, 5)
+    if random_value > 3:
+        logging.error('Dummy fail! Do not panic, just restart and try again the chance of this fail is ~20%')
+        client.update_firmware_info(state = TBFirmwareState.FAILED, error = "Dummy fail! Do not panic, just restart and try again the chance of this fail is ~20%")
+    else:
+        logging.info("Successfully updated!")
+        client.update_firmware_info(version = version_to, state = TBFirmwareState.UPDATED)
+
 
 def main():
     client = TBDeviceMqttClient("127.0.0.1", username="A2_TEST_TOKEN")
@@ -26,7 +42,7 @@ def main():
     client.get_firmware_update()
 
     # Waiting for firmware to be delivered
-    while not client.current_firmware_info[FW_STATE_ATTR] == 'UPDATED':
+    while not client.current_firmware_info[FW_STATE_ATTR] == TBFirmwareState.UPDATED.value:
         time.sleep(1)
 
     client.disconnect()

sdk_utils.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from random import randint
 from zlib import crc32
 from hashlib import sha256, sha384, sha512, md5
 import logging
@@ -78,8 +77,4 @@ def verify_checksum(firmware_data, checksum_alg, checksum):
     else:
         log.error('Client error. Unsupported checksum algorithm.')
     log.debug(checksum_of_received_firmware)
-    random_value = randint(0, 5)
-    if random_value > 3:
-        log.debug('Dummy fail! Do not panic, just restart and try again the chance of this fail is ~20%')
-        return False
     return checksum_of_received_firmware == checksum

tb_device_mqtt.py

@@ -67,6 +67,7 @@ def check_tb_paho_mqtt_installed():
 FW_CHECKSUM_ALG_ATTR = "fw_checksum_algorithm"
 FW_SIZE_ATTR = "fw_size"
 FW_STATE_ATTR = "fw_state"
+FW_ERROR_ATTR = "fw_error"
 
 REQUIRED_SHARED_KEYS = f"{FW_CHECKSUM_ATTR},{FW_CHECKSUM_ALG_ATTR},{FW_SIZE_ATTR},{FW_TITLE_ATTR},{FW_VERSION_ATTR}"
 
@@ -106,6 +107,14 @@ class TBSendMethod(Enum):
     PUBLISH = 1
     UNSUBSCRIBE = 2
 
+class TBFirmwareState(Enum):
+        IDLE = "IDLE"
+        DOWNLOADING = "DOWNLOADING"
+        DOWNLOADED = "DOWNLOADED"
+        VERIFIED =  "VERIFIED"
+        UPDATING = "UPDATING"
+        UPDATED = "UPDATED"
+        FAILED = "FAILED"
 
 class TBPublishInfo:
     TB_ERR_AGAIN = -1
@@ -474,7 +483,7 @@ class TBDeviceMqttClient:
     def __init__(self, host, port=1883, username=None, password=None, quality_of_service=None, client_id="",
                  chunk_size=0, messages_rate_limit="DEFAULT_MESSAGES_RATE_LIMIT",
                  telemetry_rate_limit="DEFAULT_TELEMETRY_RATE_LIMIT",
-                 telemetry_dp_rate_limit="DEFAULT_TELEMETRY_DP_RATE_LIMIT", max_payload_size=8196, **kwargs):
+                 telemetry_dp_rate_limit="DEFAULT_TELEMETRY_DP_RATE_LIMIT", max_payload_size=8196, on_firmware_received=None, **kwargs):
         # Added for compatibility with old versions
         if kwargs.get('rate_limit') is not None or kwargs.get('dp_rate_limit') is not None:
             messages_rate_limit = messages_rate_limit if kwargs.get('rate_limit') == "DEFAULT_RATE_LIMIT" else kwargs.get('rate_limit', messages_rate_limit) # noqa
@@ -523,37 +532,50 @@ def __init__(self, host, port=1883, username=None, password=None, quality_of_ser
         self.current_firmware_info = {
             "current_" + FW_TITLE_ATTR: "Initial",
             "current_" + FW_VERSION_ATTR: "v0",
-            FW_STATE_ATTR: "IDLE"
+            FW_STATE_ATTR: TBFirmwareState.IDLE.value,
+            FW_ERROR_ATTR: ""
         }
         self.__request_id = 0
         self.__firmware_request_id = 0
         self.__chunk_size = chunk_size
         self.firmware_received = False
+        self.set_on_firmware_received_function(on_firmware_received)
         self.rate_limits_received = False
         self.__request_service_configuration_required = False
         self.__service_loop = Thread(target=self.__service_loop, name="Service loop", daemon=True)
         self.__service_loop.start()
         self.__messages_limit_reached_set_time = (0,0)
         self.__datapoints_limit_reached_set_time = (0,0)
 
+    def update_firmware_info(self, title = None, version = None, state: TBFirmwareState = None, error = None):
+        if title is not None:
+            self.current_firmware_info["current_" + FW_TITLE_ATTR] = title
+        
+        if version is not None:
+            self.current_firmware_info["current_" + FW_VERSION_ATTR] = version
+
+        if state is not None:
+            self.current_firmware_info[FW_STATE_ATTR] = state.value
+            if state is TBFirmwareState.FAILED and error is not None:
+                self.current_firmware_info[FW_ERROR_ATTR] = error
+            else:
+                self.current_firmware_info[FW_ERROR_ATTR] = ""
+
+        self.send_telemetry(self.current_firmware_info)
+
+    def set_on_firmware_received_function(self, on_firmware_received):
+        if on_firmware_received is not None:
+            self.__on_firmware_received = on_firmware_received 
+        else:
+            self.__on_firmware_received = self.__on_firmware_received_default
+
     def __service_loop(self):
         while not self.stopped:
             if self.__request_service_configuration_required:
                 self.request_service_configuration(self.service_configuration_callback)
                 self.__request_service_configuration_required = False
             elif self.firmware_received:
-                self.current_firmware_info[FW_STATE_ATTR] = "UPDATING"
-                self.send_telemetry(self.current_firmware_info)
-                sleep(1)
-
-                self.__on_firmware_received(self.firmware_info.get(FW_VERSION_ATTR))
-
-                self.current_firmware_info = {
-                    "current_" + FW_TITLE_ATTR: self.firmware_info.get(FW_TITLE_ATTR),
-                    "current_" + FW_VERSION_ATTR: self.firmware_info.get(FW_VERSION_ATTR),
-                    FW_STATE_ATTR: "UPDATED"
-                }
-                self.send_telemetry(self.current_firmware_info)
+                self.__on_firmware_received(self, self.firmware_data, self.firmware_info.get(FW_VERSION_ATTR))
                 self.firmware_received = False
             sleep(0.05)
 
@@ -758,8 +780,7 @@ def _on_decoded_message(self, content, message):
                 self.firmware_data = b''
                 self.__current_chunk = 0
 
-                self.current_firmware_info[FW_STATE_ATTR] = "DOWNLOADING"
-                self.send_telemetry(self.current_firmware_info)
+                self.update_firmware_info(state = TBFirmwareState.DOWNLOADING)
                 sleep(1)
 
                 self.__firmware_request_id = self.__firmware_request_id + 1
@@ -769,22 +790,19 @@ def _on_decoded_message(self, content, message):
                 self.__get_firmware()
 
     def __process_firmware(self):
-        self.current_firmware_info[FW_STATE_ATTR] = "DOWNLOADED"
-        self.send_telemetry(self.current_firmware_info)
+        self.update_firmware_info(state = TBFirmwareState.DOWNLOADED)
         sleep(1)
 
         verification_result = verify_checksum(self.firmware_data, self.firmware_info.get(FW_CHECKSUM_ALG_ATTR),
                                               self.firmware_info.get(FW_CHECKSUM_ATTR))
 
         if verification_result:
             log.debug('Checksum verified!')
-            self.current_firmware_info[FW_STATE_ATTR] = "VERIFIED"
-            self.send_telemetry(self.current_firmware_info)
+            self.update_firmware_info(state = TBFirmwareState.VERIFIED)
             sleep(1)
         else:
             log.debug('Checksum verification failed!')
-            self.current_firmware_info[FW_STATE_ATTR] = "FAILED"
-            self.send_telemetry(self.current_firmware_info)
+            self.update_firmware_info(state = TBFirmwareState.FAILED, error = "Checksum verification failed!")
             self.__request_firmware_info()
             return
         self.firmware_received = True
@@ -796,10 +814,19 @@ def __get_firmware(self):
             f"v2/fw/request/{self.__firmware_request_id}/chunk/{self.__current_chunk}",
             payload=payload, qos=1)
 
-    def __on_firmware_received(self, version_to):
+    def __on_firmware_received_default(self, client, firmware_data, version_to):
+        self.update_firmware_info(state = TBFirmwareState.UPDATING)
+        sleep(1)
+
         with open(self.firmware_info.get(FW_TITLE_ATTR), "wb") as firmware_file:
-            firmware_file.write(self.firmware_data)
-        log.info('Firmware is updated!\n Current firmware version is: %s' % version_to)
+            firmware_file.write(firmware_data)
+        log.warning(f"Firmware was received and stored under {self.firmware_info.get(FW_TITLE_ATTR)}," 
+                        "but 'on_firmware_received' callback is not defined. No OTA update applied." 
+                        "Call 'set_on_firmware_received_function' to handle properly firmware update.")
+
+        self.update_firmware_info(title = self.firmware_info.get(FW_TITLE_ATTR),
+                             version = self.firmware_info.get(FW_VERSION_ATTR),
+                             state = TBFirmwareState.UPDATED)
 
     @staticmethod
     def _decode(message):