pybricksdev.dfu: add support for new SPIKE Prime hardware

Add a check for the bcdDevice when flashing SPIKE Prime, and use the
bootloader size to calculate the firmware address/size.

Author: dlech

pybricksdev/dfu.py

@@ -24,8 +24,10 @@
     SPIKE_PRIME_DFU_USB_PID,
 )
 
-FIRMWARE_ADDRESS = 0x08008000
-FIRMWARE_SIZE = 1 * 1024 * 1024 - 32 * 1024  # 1MiB - 32KiB
+BOOTLOADER_SIZE_32K = 32 * 1024
+BOOTLOADER_SIZE_64K = 64 * 1024
+FLASH_BASE_ADDRESS = 0x08000000
+FLASH_SIZE = 1 * 1024 * 1024
 
 
 ALL_PIDS = {
@@ -36,6 +38,20 @@
 ALL_DEVICES = [f"{LEGO_USB_VID:04x}:{pid:04x}" for pid in ALL_PIDS.keys()]
 
 
+def _get_bootloader_size(pid: int, bcd_device: int | None) -> int:
+    """Gets bootloader size for the connected DFU device."""
+    # New hardware revision of SPIKE Prime released in 2026 has a larger bootloader.
+    if pid == SPIKE_PRIME_DFU_USB_PID and bcd_device == 0x0300:
+        return BOOTLOADER_SIZE_64K
+
+    return BOOTLOADER_SIZE_32K
+
+
+def _get_firmware_region(bootloader_size: int) -> tuple[int, int]:
+    """Gets firmware flash address and size from bootloader size."""
+    return FLASH_BASE_ADDRESS + bootloader_size, FLASH_SIZE - bootloader_size
+
+
 def _get_dfu_util() -> ContextManager[os.PathLike[str] | str]:
     """Gets ``dfu-util`` command line tool path.
 
@@ -62,9 +78,9 @@ def _get_dfu_util() -> ContextManager[os.PathLike[str] | str]:
     return nullcontext(dfu_util)
 
 
-def _get_vid_pid(dfu_util: os.PathLike[str] | str) -> str:
+def _get_dfu_device_info(dfu_util: os.PathLike[str] | str) -> tuple[str, int]:
     """
-    Gets the VID and PID of a connected LEGO DFU device.
+    Gets the VID:PID and bootloader size of a connected LEGO DFU device.
 
     Returns: The first matching LEGO DFU device from ``dfu-util --list``
 
@@ -79,7 +95,16 @@ def _get_vid_pid(dfu_util: os.PathLike[str] | str) -> str:
         dev_id = line[line.index(b"[") + 1 : line.index(b"]")].decode()
 
         if dev_id in ALL_DEVICES:
-            return dev_id
+            pid = int(dev_id.split(":", maxsplit=1)[1], 16)
+            bcd_device = None
+
+            try:
+                i = line.index(b"ver=") + 4
+                bcd_device = int(line[i : i + 4].decode(), 16)
+            except (ValueError, IndexError):
+                pass
+
+            return dev_id, _get_bootloader_size(pid, bcd_device)
 
     raise RuntimeError("No LEGO DFU USB device found")
 
@@ -98,6 +123,8 @@ def backup_dfu(file: BinaryIO) -> None:
         # if libusb was not found, try using dfu-util command line tool
 
         with _get_dfu_util() as dfu_util:
+            dev_id, bootloader_size = _get_dfu_device_info(dfu_util)
+            firmware_address, firmware_size = _get_firmware_region(bootloader_size)
             file.close()
 
             # dfu-util won't overwrite existing files so we have to do that first
@@ -108,11 +135,11 @@ def backup_dfu(file: BinaryIO) -> None:
                     [
                         dfu_util,
                         "--device",
-                        f",{_get_vid_pid(dfu_util)}",
+                        f",{dev_id}",
                         "--alt",
                         "0",
                         "--dfuse-address",
-                        f"{FIRMWARE_ADDRESS}:{FIRMWARE_SIZE}",
+                        f"{firmware_address}:{firmware_size}",
                         "--upload",
                         file.name,
                     ]
@@ -140,18 +167,20 @@ def restore_dfu(file: BinaryIO) -> None:
         # if libusb was not found, try using dfu-util command line tool
 
         with _get_dfu_util() as dfu_util:
+            dev_id, bootloader_size = _get_dfu_device_info(dfu_util)
+            firmware_address, _ = _get_firmware_region(bootloader_size)
             file.close()
 
             exit(
                 call(
                     [
                         dfu_util,
                         "--device",
-                        f",{_get_vid_pid(dfu_util)}",
+                        f",{dev_id}",
                         "--alt",
                         "0",
                         "--dfuse-address",
-                        f"{FIRMWARE_ADDRESS}:leave",
+                        f"{firmware_address}:leave",
                         "--download",
                         file.name,
                     ]
@@ -164,7 +193,6 @@ def flash_dfu(firmware_bin: bytes, metadata: AnyFirmwareMetadata) -> None:
 
     with TemporaryDirectory() as out_dir:
         outfile = os.path.join(out_dir, "firmware.dfu")
-        target: dfu_create.Image = {"address": FIRMWARE_ADDRESS, "data": firmware_bin}
 
         try:
             # Determine correct product ID
@@ -178,7 +206,8 @@ def flash_dfu(firmware_bin: bytes, metadata: AnyFirmwareMetadata) -> None:
                 )
                 exit(1)
 
-            product_id = devices[0].idProduct
+            product_id = int(devices[0].idProduct)
+            bcd_device = int(devices[0].bcdDevice)
             if product_id not in ALL_PIDS:
                 print(f"Unknown USB product ID: {product_id:04X}", file=sys.stderr)
                 exit(1)
@@ -187,6 +216,13 @@ def flash_dfu(firmware_bin: bytes, metadata: AnyFirmwareMetadata) -> None:
                 print("Incorrect firmware type for this hub", file=sys.stderr)
                 exit(1)
 
+            bootloader_size = _get_bootloader_size(product_id, bcd_device)
+            firmware_address, _ = _get_firmware_region(bootloader_size)
+            target: dfu_create.Image = {
+                "address": firmware_address,
+                "data": firmware_bin,
+            }
+
             # Create dfu file
             device = "0x{0:04x}:0x{1:04x}".format(LEGO_USB_VID, product_id)
             dfu_create.build(outfile, [[target]], device)
@@ -225,6 +261,9 @@ def flash_dfu(firmware_bin: bytes, metadata: AnyFirmwareMetadata) -> None:
             # if libusb was not found, try using dfu-util command line tool
 
             with _get_dfu_util() as dfu_util:
+                dev_id, bootloader_size = _get_dfu_device_info(dfu_util)
+                firmware_address, _ = _get_firmware_region(bootloader_size)
+
                 with open(os.path.join(out_dir, "firmware.bin"), "wb") as bin_file:
                     bin_file.write(firmware_bin)
 
@@ -233,14 +272,14 @@ def flash_dfu(firmware_bin: bytes, metadata: AnyFirmwareMetadata) -> None:
                         [
                             dfu_util,
                             "--device",
-                            f",{_get_vid_pid(dfu_util)}",
+                            f",{dev_id}",
                             "--alt",
                             "0",
                             # We have to use dfuse option to be able to use
                             # "leave" to exit DFU mode after flashing. --reset
                             # doesn't work on Windows, so we can't use a .dfu file
                             "--dfuse-address",
-                            f"{target['address']}:leave",
+                            f"{firmware_address}:leave",
                             "--download",
                             bin_file.name,
                         ]