tools: add script to generate DFU binaries for ESP32-S2

tools/mkdfu.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python
+#
+# This program creates archives compatible with ESP32-S* ROM DFU implementation.
+#
+# The archives are in CPIO format. Each file which needs to be flashed is added to the archive
+# as a separate file. In addition to that, a special index file, 'dfuinfo0.dat', is created.
+# This file must be the first one in the archive. It contains binary structures describing each
+# subsequent file (for example, where the file needs to be flashed/loaded).
+#
+
+import argparse
+from collections import namedtuple
+import hashlib
+import os
+import struct
+import zlib
+
+try:
+    import typing
+except ImportError:
+    # Only used for type annotations
+    pass
+
+# CPIO ("new ASCII") format related things
+CPIO_MAGIC = b"070701"
+CPIO_STRUCT = b"=6s" + b"8s" * 13
+CPIOHeader = namedtuple(
+    "CPIOHeader",
+    [
+        "magic",
+        "ino",
+        "mode",
+        "uid",
+        "gid",
+        "nlink",
+        "mtime",
+        "filesize",
+        "devmajor",
+        "devminor",
+        "rdevmajor",
+        "rdevminor",
+        "namesize",
+        "check",
+    ],
+)
+CPIO_TRAILER = "TRAILER!!!"
+
+
+def make_cpio_header(
+    filename_len, file_len, is_trailer=False
+):  # type: (int, int, bool) -> CPIOHeader
+    """ Returns CPIOHeader for the given file name and file size """
+
+    def as_hex(val):  # type: (int) -> bytes
+        return "{:08x}".format(val).encode("ascii")
+
+    hex_0 = as_hex(0)
+    mode = hex_0 if is_trailer else as_hex(0o0100644)
+    nlink = as_hex(1) if is_trailer else hex_0
+    return CPIOHeader(
+        magic=CPIO_MAGIC,
+        ino=hex_0,
+        mode=mode,
+        uid=hex_0,
+        gid=hex_0,
+        nlink=nlink,
+        mtime=hex_0,
+        filesize=as_hex(file_len),
+        devmajor=hex_0,
+        devminor=hex_0,
+        rdevmajor=hex_0,
+        rdevminor=hex_0,
+        namesize=as_hex(filename_len),
+        check=hex_0,
+    )
+
+
+# DFU format related things
+# Structure of one entry in dfuinfo0.dat
+DFUINFO_STRUCT = b"<I I 64s 16s"
+DFUInfo = namedtuple("DFUInfo", ["address", "flags", "name", "md5"])
+DFUINFO_FILE = "dfuinfo0.dat"
+# Structure which gets added at the end of the entire DFU file
+DFUSUFFIX_STRUCT = b"<H H H H 3s B"
+DFUSuffix = namedtuple(
+    "DFUSuffix", ["bcd_device", "pid", "vid", "bcd_dfu", "sig", "len"]
+)
+ESPRESSIF_VID = 12346
+# TODO: set PID based on the chip type (add a command line argument)
+DFUSUFFIX_DEFAULT = DFUSuffix(0xFFFF, 0xFFFF, ESPRESSIF_VID, 0x0100, b"UFD", 16)
+# This CRC32 gets added after DFUSUFFIX_STRUCT
+DFUCRC_STRUCT = b"<I"
+
+
+def dfu_crc(data, crc=0):  # type: (bytes, int) -> int
+    """ Calculate CRC32/JAMCRC of data, with an optional initial value """
+    uint32_max = 0xFFFFFFFF
+    return uint32_max - (zlib.crc32(data, crc) & uint32_max)
+
+
+def pad_bytes(b, multiple, padding=b"\x00"):  # type: (bytes, int, bytes) -> bytes
+    """ Pad 'b' to a length divisible by 'multiple' """
+    padded_len = (len(b) + multiple - 1) // multiple * multiple
+    return b + padding * (padded_len - len(b))
+
+
+class EspDfuWriter(object):
+    def __init__(self, dest_file):  # type: (typing.BinaryIO) -> None
+        self.dest = dest_file
+        self.entries = []  # type: typing.List[bytes]
+        self.index = []  # type: typing.List[DFUInfo]
+
+    def add_file(self, flash_addr, path):  # type: (int, str) -> None
+        """ Add file to be written into flash at given address """
+        with open(path, "rb") as f:
+            self._add_cpio_flash_entry(os.path.basename(path), flash_addr, f.read())
+
+    def finish(self):  # type: () -> None
+        """ Write DFU file """
+        # Prepare and add dfuinfo0.dat file
+        dfuinfo = b"".join([struct.pack(DFUINFO_STRUCT, *item) for item in self.index])
+        self._add_cpio_entry(DFUINFO_FILE, dfuinfo, first=True)
+
+        # Add CPIO archive trailer
+        self._add_cpio_entry(CPIO_TRAILER, b"", trailer=True)
+
+        # Combine all the entries and pad the file
+        out_data = b"".join(self.entries)
+        cpio_block_size = 10240
+        out_data = pad_bytes(out_data, cpio_block_size)
+
+        # Add DFU suffix and CRC
+        out_data += struct.pack(DFUSUFFIX_STRUCT, *DFUSUFFIX_DEFAULT)
+        out_data += struct.pack(DFUCRC_STRUCT, dfu_crc(out_data))
+
+        # Finally write the entire binary
+        self.dest.write(out_data)
+
+    def _add_cpio_flash_entry(
+        self, filename, flash_addr, data
+    ):  # type: (str, int, bytes) -> None
+        md5 = hashlib.md5()
+        md5.update(data)
+        self.index.append(
+            DFUInfo(
+                address=flash_addr,
+                flags=0,
+                name=filename.encode("utf-8"),
+                md5=md5.digest(),
+            )
+        )
+        self._add_cpio_entry(filename, data)
+
+    def _add_cpio_entry(
+        self, filename, data, first=False, trailer=False
+    ):  # type: (str, bytes, bool, bool) -> None
+        filename_b = filename.encode("utf-8") + b"\x00"
+        cpio_header = make_cpio_header(len(filename_b), len(data), is_trailer=trailer)
+        entry = pad_bytes(
+            struct.pack(CPIO_STRUCT, *cpio_header) + filename_b, 4
+        ) + pad_bytes(data, 4)
+        if not first:
+            self.entries.append(entry)
+        else:
+            self.entries.insert(0, entry)
+
+
+def action_write(args):
+    writer = EspDfuWriter(args.output_file)
+    for addr, file in args.files:
+        writer.add_file(addr, file)
+    writer.finish()
+
+
+class WriteArgsAction(argparse.Action):
+    """ Helper for argparse to parse <address, filename> argument pairs """
+
+    def __init__(self, *args, **kwargs):
+        super(WriteArgsAction, self).__init__(*args, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        # TODO: add validation
+        addr = 0
+        result = []
+        for i, value in enumerate(values):
+            if i % 2 == 0:
+                addr = int(value, 0)
+            else:
+                result.append((addr, value))
+
+        setattr(namespace, self.dest, result)
+
+
+def main():
+    parser = argparse.ArgumentParser("mkdfu")
+
+    # Provision to add "info" command
+    subparsers = parser.add_subparsers(dest="command")
+    write_parser = subparsers.add_parser("write")
+    write_parser.add_argument("-o", "--output-file", type=argparse.FileType("wb"))
+    write_parser.add_argument(
+        "files", metavar="<address> <file>", action=WriteArgsAction, nargs="+"
+    )
+
+    args = parser.parse_args()
+    print(repr(args))
+    if args.command == "write":
+        action_write(args)
+
+
+if __name__ == "__main__":
+    main()