dfu: fix dfu-util --device args

It seems that only the first arg after the first "," in the dfu-util
--device arg is used to match attached devices. Since SPIKE Prime and
MINDSTORMS RI have different product IDs, we need to get the right ID
to make it work for both hubs. We can do this by scraping the ID from
the dfu-util --list command.

Tested with dfu-util 0.9 on Ubuntu 20.04.

Issue: pybricks/support#420

Author: dlech

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Fixed crash in CRC32 checksum.
+- Fixed flashing with `dfu-util` not always working ([support#420]).
 
 ## [1.0.0-alpha.12]
 
@@ -105,6 +106,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Typo in `pip` arguments `README.md`.
 
 
+<!-- let's try to keep these sorted alphabetically -->
+[support#420]: https://github.com/pybricks/support/issues/420
+
 [Unreleased]: https://github.com/pybricks/pybricksdev/compare/v1.0.0-alpha.12..HEAD
 [1.0.0-alpha.12]: https://github.com/pybricks/pybricksdev/compare/v1.0.0-alpha.11...v1.0.0-alpha.12
 [1.0.0-alpha.11]: https://github.com/pybricks/pybricksdev/compare/v1.0.0-alpha.10...v1.0.0-alpha.11

pybricksdev/_dfu_create.py

@@ -104,7 +104,7 @@ def build(file, targets, device=DEFAULT_DEVICE):
         )
         data += tdata
     data = struct.pack("<5sBIB", b"DfuSe", 1, len(data) + 11, len(targets)) + data
-    v, d = map(lambda x: int(x, 0) & 0xFFFF, device.split(":", 1))
+    v, d = map(lambda x: int(x, 16) & 0xFFFF, device.split(":", 1))
     data += struct.pack("<4H3sB", 0, d, v, 0x011A, b"UFD", 16)
     crc = compute_crc(data)
     data += struct.pack("<I", crc)

pybricksdev/dfu.py

@@ -8,7 +8,7 @@
 
 from contextlib import nullcontext
 from importlib.resources import path
-from subprocess import DEVNULL, call, check_call
+from subprocess import DEVNULL, PIPE, call, check_call, run
 from tempfile import TemporaryDirectory
 from typing import BinaryIO, ContextManager
 
@@ -24,8 +24,9 @@
 MINDSTORMS_INVENTOR_PID = 0x0011
 
 
-SPIKE_PRIME_DEVICE = f"0x{LEGO_VID:04x}:0x{SPIKE_PRIME_PID:04x}"
-MINDSTORMS_INVENTOR_DEVICE = f"0x{LEGO_VID:04x}:0x{MINDSTORMS_INVENTOR_PID:04x}"
+SPIKE_PRIME_DEVICE = f"{LEGO_VID:04x}:{SPIKE_PRIME_PID:04x}"
+MINDSTORMS_INVENTOR_DEVICE = f"{LEGO_VID:04x}:{MINDSTORMS_INVENTOR_PID:04x}"
+ALL_DEVICES = [SPIKE_PRIME_DEVICE, MINDSTORMS_INVENTOR_DEVICE]
 
 
 def _get_dfu_util() -> ContextManager[os.PathLike]:
@@ -54,6 +55,28 @@ def _get_dfu_util() -> ContextManager[os.PathLike]:
     return nullcontext(dfu_util)
 
 
+def _get_vid_pid(dfu_util: os.PathLike) -> str:
+    """
+    Gets the VID and PID of a connected LEGO DFU device.
+
+    Returns: The first matching LEGO DFU device from ``dfu-util --list``
+
+    Raises: RuntimeError: No matching hubs found.
+    """
+    proc = run([dfu_util, "--list"], stdout=PIPE, check=True)
+
+    for line in proc.stdout.splitlines():
+        if not line.startswith(b"Found DFU:"):
+            continue
+
+        dev_id = line[line.index(b"[") + 1 : line.index(b"]")].decode()
+
+        if dev_id in ALL_DEVICES:
+            return dev_id
+
+    raise RuntimeError("No LEGO DFU USB device found")
+
+
 def backup_dfu(file: BinaryIO) -> None:
     """Backs up device data via DFU.
 
@@ -79,7 +102,7 @@ def backup_dfu(file: BinaryIO) -> None:
                     [
                         dfu_util,
                         "--device",
-                        f",{SPIKE_PRIME_DEVICE},{MINDSTORMS_INVENTOR_DEVICE}",
+                        f",{_get_vid_pid(dfu_util)}",
                         "--alt",
                         "0",
                         "--dfuse-address",
@@ -119,7 +142,7 @@ def restore_dfu(file: BinaryIO) -> None:
                     [
                         dfu_util,
                         "--device",
-                        f",{SPIKE_PRIME_DEVICE},{MINDSTORMS_INVENTOR_DEVICE}",
+                        f",{_get_vid_pid(dfu_util)}",
                         "--alt",
                         "0",
                         "--dfuse-address",
@@ -197,16 +220,16 @@ def flash_dfu(firmware_bin: bytes, metadata: dict) -> None:
 
             with _get_dfu_util() as dfu_util:
 
-                # Exact device product ID doesn't matter here since we are using
-                # the --device command line option below.
-                _dfu_create.build(outfile, [[target]], SPIKE_PRIME_DEVICE)
+                dev_id = _get_vid_pid(dfu_util)
+
+                _dfu_create.build(outfile, [[target]], dev_id)
 
                 exit(
                     call(
                         [
                             dfu_util,
                             "--device",
-                            f",{SPIKE_PRIME_DEVICE},{MINDSTORMS_INVENTOR_DEVICE}",
+                            f",{dev_id}",
                             "--alt",
                             "0",
                             "--download",