From c06d2393111b5ea13fc762fcf4f641681c78e671 Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Sun, 26 Nov 2023 17:25:41 -0600
Subject: [PATCH] pybricksdev.usb: add new usb subpackage

We were starting to get repeated defines of USB stuff, so let's put it
all in one place.
---
 pybricksdev/connections/ev3.py  |  9 ++-------
 pybricksdev/connections/lego.py |  6 +++++-
 pybricksdev/dfu.py              | 19 ++++++++-----------
 pybricksdev/usb/__init__.py     | 23 +++++++++++++++++++++++
 4 files changed, 38 insertions(+), 19 deletions(-)
 create mode 100644 pybricksdev/usb/__init__.py

diff --git a/pybricksdev/connections/ev3.py b/pybricksdev/connections/ev3.py
index 627292e..9346cb5 100644
--- a/pybricksdev/connections/ev3.py
+++ b/pybricksdev/connections/ev3.py
@@ -10,10 +10,7 @@
 import hid
 
 from ..tools import chunk
-
-LEGO_VENDOR_ID = 0x0694
-EV3_PRODUCT_ID = 0x0005
-EV3_BOOTLOADER_PRODUCT_ID = 0x0006
+from ..usb import LEGO_USB_VID, LegoUsbPid
 
 
 class MessageType(enum.IntEnum):
@@ -70,9 +67,7 @@ def open(self) -> None:
         """
         Opens an HID connection to the EV3 bootloader.
         """
-        self._device.open(
-            vendor_id=LEGO_VENDOR_ID, product_id=EV3_BOOTLOADER_PRODUCT_ID
-        )
+        self._device.open(vendor_id=LEGO_USB_VID, product_id=LegoUsbPid.EV3_FW_UPDATE)
 
     def close(self) -> None:
         """
diff --git a/pybricksdev/connections/lego.py b/pybricksdev/connections/lego.py
index 079f777..d9349c9 100644
--- a/pybricksdev/connections/lego.py
+++ b/pybricksdev/connections/lego.py
@@ -8,6 +8,7 @@
 from serial.tools import list_ports
 
 from ..tools import chunk
+from ..usb import LEGO_USB_VID
 from .pybricks import PybricksHub
 
 FILE_PACKET_SIZE = 1024
@@ -77,7 +78,10 @@ async def connect(self, device=None):
         port = None
         devices = list_ports.comports()
         for dev in devices:
-            if dev.product == "LEGO Technic Large Hub in FS Mode" or dev.vid == 0x0694:
+            if (
+                dev.product == "LEGO Technic Large Hub in FS Mode"
+                or dev.vid == LEGO_USB_VID
+            ):
                 port = dev.device
                 break
 
diff --git a/pybricksdev/dfu.py b/pybricksdev/dfu.py
index c836c41..99f00de 100644
--- a/pybricksdev/dfu.py
+++ b/pybricksdev/dfu.py
@@ -16,20 +16,17 @@
 from . import resources
 from ._vendored import dfu_create, dfu_upload
 from .ble.lwp3.bytecodes import HubKind
+from .usb import LEGO_USB_VID, LegoUsbPid
 
 FIRMWARE_ADDRESS = 0x08008000
 FIRMWARE_SIZE = 1 * 1024 * 1024 - 32 * 1024  # 1MiB - 32KiB
-LEGO_VID = 0x0694
-SPIKE_PRIME_PID = 0x0008
-SPIKE_ESSENTIAL_PID = 0x000C
-MINDSTORMS_INVENTOR_PID = 0x0011
 
 ALL_PIDS = {
-    MINDSTORMS_INVENTOR_PID: HubKind.TECHNIC_LARGE,
-    SPIKE_ESSENTIAL_PID: HubKind.TECHNIC_SMALL,
-    SPIKE_PRIME_PID: HubKind.TECHNIC_LARGE,
+    LegoUsbPid.SPIKE_PRIME_DFU: HubKind.TECHNIC_LARGE,
+    LegoUsbPid.SPIKE_ESSENTIAL_DFU: HubKind.TECHNIC_SMALL,
+    LegoUsbPid.ROBOT_INVENTOR_DFU: HubKind.TECHNIC_LARGE,
 }
-ALL_DEVICES = [f"{LEGO_VID:04x}:{pid:04x}" for pid in ALL_PIDS.keys()]
+ALL_DEVICES = [f"{LEGO_USB_VID:04x}:{pid:04x}" for pid in ALL_PIDS.keys()]
 
 
 def _get_dfu_util() -> ContextManager[os.PathLike]:
@@ -165,7 +162,7 @@ def flash_dfu(firmware_bin: bytes, metadata: dict) -> None:
         try:
             # Determine correct product ID
 
-            devices = dfu_upload.get_dfu_devices(idVendor=LEGO_VID)
+            devices = dfu_upload.get_dfu_devices(idVendor=LEGO_USB_VID)
             if not devices:
                 print(
                     "No DFU devices found.",
@@ -184,11 +181,11 @@ def flash_dfu(firmware_bin: bytes, metadata: dict) -> None:
                 exit(1)
 
             # Create dfu file
-            device = "0x{0:04x}:0x{1:04x}".format(LEGO_VID, product_id)
+            device = "0x{0:04x}:0x{1:04x}".format(LEGO_USB_VID, product_id)
             dfu_create.build(outfile, [[target]], device)
 
             # Init dfu tool
-            dfu_upload.__VID = LEGO_VID
+            dfu_upload.__VID = LEGO_USB_VID
             dfu_upload.__PID = product_id
             dfu_upload.init()
             elements = dfu_upload.read_dfu_file(outfile)
diff --git a/pybricksdev/usb/__init__.py b/pybricksdev/usb/__init__.py
new file mode 100644
index 0000000..5728e99
--- /dev/null
+++ b/pybricksdev/usb/__init__.py
@@ -0,0 +1,23 @@
+"""
+Common USB definitions.
+"""
+
+import enum as _enum
+
+LEGO_USB_VID = 0x0694
+NXT_SAMBA_VID = 0x03EB
+NXT_SAMBA_PID = 0x6124
+
+
+class LegoUsbPid(_enum.IntEnum):
+    RCX_IR_TOWER = 0x0001
+    NXT = 0x0002
+    WEDO = 0x0003
+    EV3 = 0x0005
+    EV3_FW_UPDATE = 0x0006
+    SPIKE_PRIME_DFU = 0x0008
+    SPIKE_PRIME = 0x0009
+    SPIKE_ESSENTIAL_DFU = 0x000C
+    SPIKE_ESSENTIAL = 0x000D
+    ROBOT_INVENTOR = 0x0010
+    ROBOT_INVENTOR_DFU = 0x0011