From 8ce25eb2fb2660fece3cb980779e5f8337fa0fa5 Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Sun, 26 Nov 2023 18:48:33 -0600
Subject: [PATCH] pybricksdev.usb: add Pybricks USB device
 class/subclass/protocol

This adds constants for the Pybricks USB device class, subclass and
protocol. These are used to identify hubs running the Pybricks firmware
when connected via USB.

The `usb` command of the CLI is modified to only accept Pybricks hubs
now. If we still want to support running programs on hubs running
official LEGO firmware, we can add that back with a different command
name.
---
 pybricksdev/cli/__init__.py | 25 ++++++++++++++++++-------
 pybricksdev/usb/__init__.py |  5 +++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py
index 5d9d2d6..86462bc 100644
--- a/pybricksdev/cli/__init__.py
+++ b/pybricksdev/cli/__init__.py
@@ -20,6 +20,12 @@
 
 from .. import __name__ as MODULE_NAME
 from .. import __version__ as MODULE_VERSION
+from ..usb import (
+    LEGO_USB_VID,
+    PYBRICKS_USB_DEVICE_CLASS,
+    PYBRICKS_USB_DEVICE_PROTOCOL,
+    PYBRICKS_USB_DEVICE_SUBCLASS,
+)
 
 PROG_NAME = (
     f"{path.basename(sys.executable)} -m {MODULE_NAME}"
@@ -171,11 +177,11 @@ def add_parser(self, subparsers: argparse._SubParsersAction):
             )
 
     async def run(self, args: argparse.Namespace):
+        from usb.core import find as find_usb
+
         from ..ble import find_device as find_ble
         from ..connections.ev3dev import EV3Connection
-        from ..connections.lego import REPLHub
         from ..connections.pybricks import PybricksHubBLE, PybricksHubUSB
-        from usb.core import find as find_usb
 
         # Pick the right connection
         if args.conntype == "ssh":
@@ -194,12 +200,17 @@ async def run(self, args: argparse.Namespace):
             hub = PybricksHubBLE(device_or_address)
 
         elif args.conntype == "usb":
-            device_or_address = find_usb(idVendor=0x0483, idProduct=0x5740)
+            usb_device = find_usb(
+                idVendor=LEGO_USB_VID,
+                bDeviceClass=PYBRICKS_USB_DEVICE_CLASS,
+                bDeviceSubClass=PYBRICKS_USB_DEVICE_SUBCLASS,
+                bDeviceProtocol=PYBRICKS_USB_DEVICE_PROTOCOL,
+            )
+            if usb_device is None:
+                print("No Pybricks USB device found.", file=sys.stderr)
+                exit(1)
 
-            if device_or_address is not None and device_or_address.product == "Pybricks Hub":
-                hub = PybricksHubUSB(device_or_address)
-            else:
-                hub = REPLHub()
+            hub = PybricksHubUSB(usb_device)
         else:
             raise ValueError(f"Unknown connection type: {args.conntype}")
 
diff --git a/pybricksdev/usb/__init__.py b/pybricksdev/usb/__init__.py
index 5728e99..3ad7665 100644
--- a/pybricksdev/usb/__init__.py
+++ b/pybricksdev/usb/__init__.py
@@ -21,3 +21,8 @@ class LegoUsbPid(_enum.IntEnum):
     SPIKE_ESSENTIAL = 0x000D
     ROBOT_INVENTOR = 0x0010
     ROBOT_INVENTOR_DFU = 0x0011
+
+
+PYBRICKS_USB_DEVICE_CLASS = 0xFF
+PYBRICKS_USB_DEVICE_SUBCLASS = 0xC5
+PYBRICKS_USB_DEVICE_PROTOCOL = 0xF5