micropython/espflash: Use hashlib.md5 to generate file digest.

iabdalkader · iabdalkader · commit 1918c47035fa · 2022-11-03T22:39:34.000+01:00
* Use hashlib to generate the file MD5 digest when available,
and if none is given in the args.
diff --git a/micropython/espflash/espflash.py b/micropython/espflash/espflash.py
@@ -32,6 +32,7 @@
 from machine import UART
 from micropython import const
 from time import sleep
+import binascii
 
 _CMD_SYNC = const(0x08)
 _CMD_CHANGE_BAUDRATE = const(0x0F)
@@ -76,6 +77,13 @@ def __init__(
         self.reset_pin = Pin(reset, Pin.OUT)
         self.gpio0_pin = Pin(gpio0, Pin.OUT)
         self.set_baudrate(self.uart_baudrate)
+        self.md5sum = None
+        try:
+            import hashlib
+
+            self.md5sum = hashlib.md5()
+        except ImportError:
+            pass
 
     def _log(self, data, out=True):
         if self.log:
@@ -266,8 +274,11 @@ def flash_write_file(self, path, blksize=0x1000):
             subseq = 0
             for i in range(total_blocks):
                 buf = f.read(blksize)
+                # Update digest
+                if self.md5sum is not None:
+                    self.md5sum.update(buf)
+                # The last data block should be padded to the block size with 0xFF bytes.
                 if len(buf) < blksize:
-                    # The last data block should be padded to the block size with 0xFF bytes.
                     buf += b"\xFF" * (blksize - len(buf))
                 checksum = self.checksum(buf)
                 if seq % erase_blocks == 0:
@@ -288,15 +299,22 @@ def flash_write_file(self, path, blksize=0x1000):
 
         print("Flash write finished")
 
-    def flash_verify_file(self, path, md5sum, offset=0):
+    def flash_verify_file(self, path, digest=None, offset=0):
+        if digest is None:
+            if self.md5sum is None:
+                raise Exception(f"MD5 checksum missing.")
+            digest = binascii.hexlify(self.md5sum.digest())
+
         size = os.stat(path)[6]
         val, data = self.command(_CMD_SPI_FLASH_MD5, struct.pack("<IIII", offset, size, 0, 0))
-        print(f"Flash verify file MD5 {md5sum}")
-        print(f"Flash verify flash MD5 {bytes(data[0:32])}")
-        if md5sum == data[0:32]:
-            print("Firmware write verified")
+
+        print(f"Flash verify: File  MD5 {digest}")
+        print(f"Flash verify: Flash MD5 {bytes(data[0:32])}")
+
+        if digest == data[0:32]:
+            print("Firmware verified.")
         else:
-            raise Exception(f"Firmware verification failed")
+            raise Exception(f"Firmware verification failed.")
 
     def reboot(self):
         payload = struct.pack("<I", 0)
