From 2eed9780ba7074de9e464a2bc771ad14f0332a6c Mon Sep 17 00:00:00 2001
From: David Michieli <dmm@planetinnovation.com.au>
Date: Tue, 16 Feb 2021 12:34:34 +1100
Subject: [PATCH] stm32/mboot: Add unpack-dfu command to mboot_pack_dfu.py
 tool.

This command unpacks a previously packed DFU file, writing out a DFU which
should be the same as the original (before packing).
---
 ports/stm32/mboot/mboot_pack_dfu.py | 26 ++++++++++++++++++++++++++
 tools/ci.sh                         |  5 +++++
 2 files changed, 31 insertions(+)

diff --git a/ports/stm32/mboot/mboot_pack_dfu.py b/ports/stm32/mboot/mboot_pack_dfu.py
index 47382f5910..540057e06e 100644
--- a/ports/stm32/mboot/mboot_pack_dfu.py
+++ b/ports/stm32/mboot/mboot_pack_dfu.py
@@ -211,9 +211,16 @@ def pack_dfu(keys, args):
 
 
 def verify_pack_dfu(keys, filename):
+    """Verify packed dfu file against keys. Gathers decrypted binary data."""
     full_sig = pyhy.hydro_sign(MBOOT_PACK_HYDRO_CONTEXT)
     _, elems = dfu_read(filename)
+    base_addr = None
+    binary_data = b""
+
     for addr, data in elems:
+        if base_addr is None:
+            base_addr = addr
+
         header = struct.unpack("<BBBBII", data[:12])
         chunk = data[12 : 12 + header[5]]
         sig = data[12 + header[5] :]
@@ -221,6 +228,7 @@ def verify_pack_dfu(keys, filename):
             sig, data[:12] + chunk, MBOOT_PACK_HYDRO_CONTEXT, keys.sign_pk
         )
         assert sig_pass
+
         if header[1] == MBOOT_PACK_CHUNK_FULL_SIG:
             actual_sig = chunk[-64:]
         else:
@@ -231,8 +239,21 @@ def verify_pack_dfu(keys, filename):
             if header[1] == MBOOT_PACK_CHUNK_FW_GZIP:
                 chunk = zlib.decompress(chunk, wbits=-15)
             full_sig.update(chunk)
+            assert addr == base_addr + len(binary_data)
+            binary_data += chunk
+
     full_sig_pass = full_sig.final_verify(actual_sig, keys.sign_pk)
     assert full_sig_pass
+    return [{"address": base_addr, "data": binary_data}]
+
+
+def unpack_dfu(keys, args):
+    # Load previously generated keys.
+    keys.load()
+
+    # Build a DFU file from the decrypted binary data.
+    data = verify_pack_dfu(keys, args.infile[0])
+    dfu.build(args.outfile[0], [data])
 
 
 def main():
@@ -250,6 +271,11 @@ def main():
     parser_ed.add_argument("outfile", nargs=1, help="output DFU file")
     parser_ed.set_defaults(func=pack_dfu)
 
+    parser_dd = subparsers.add_parser("unpack-dfu", help="decrypt a signed/encrypted DFU file")
+    parser_dd.add_argument("infile", nargs=1, help="input packed DFU file")
+    parser_dd.add_argument("outfile", nargs=1, help="output DFU file")
+    parser_dd.set_defaults(func=unpack_dfu)
+
     args = cmd_parser.parse_args()
 
     keys = Keys(args.keys)
diff --git a/tools/ci.sh b/tools/ci.sh
index 5e0686f301..a815e9483e 100755
--- a/tools/ci.sh
+++ b/tools/ci.sh
@@ -223,6 +223,11 @@ function ci_stm32_nucleo_build {
     make ${MAKEOPTS} -C ports/stm32 BOARD=NUCLEO_L476RG DEBUG=1
     make ${MAKEOPTS} -C ports/stm32 BOARD=NUCLEO_WB55
     make ${MAKEOPTS} -C ports/stm32/mboot BOARD=NUCLEO_WB55
+    # Test mboot_pack_dfu.py created a valid file, and that its unpack-dfu command works.
+    BOARD_WB55=ports/stm32/boards/NUCLEO_WB55
+    BUILD_WB55=ports/stm32/build-NUCLEO_WB55
+    python3 ports/stm32/mboot/mboot_pack_dfu.py -k $BOARD_WB55/mboot_keys.h unpack-dfu $BUILD_WB55/firmware.pack.dfu $BUILD_WB55/firmware.unpack.dfu
+    diff $BUILD_WB55/firmware.unpack.dfu $BUILD_WB55/firmware.dfu
 }
 
 ########################################################################################