Add support for USB connections

Adds a new transport to manage USB connections.

Signed-off-by: Nate Karstens <[email protected]>

Author: nkarstens

pybricksdev/ble/pybricks.py

@@ -328,6 +328,12 @@ def _standard_uuid(short: int) -> str:
 .. availability:: Since Pybricks protocol v1.0.0.
 """
 
+DEVICE_NAME_UUID = _standard_uuid(0x2A00)
+"""Standard Device Name UUID
+
+.. availability:: Since Pybricks protocol v1.0.0.
+"""
+
 FW_REV_UUID = _standard_uuid(0x2A26)
 """Standard Firmware Revision String characteristic UUID
 

pybricksdev/connections/pybricks.py

@@ -8,6 +8,7 @@
 import os
 import struct
 from typing import Awaitable, Callable, List, Optional, Tuple, TypeVar
+from uuid import UUID
 
 import reactivex.operators as op
 import semver
@@ -19,9 +20,15 @@
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
+from usb.control import get_descriptor
+from usb.core import Device as USBDevice
+from usb.core import Endpoint, USBTimeoutError
+from usb.util import ENDPOINT_IN, ENDPOINT_OUT, endpoint_direction, find_descriptor
+
 from ..ble.lwp3.bytecodes import HubKind
 from ..ble.nus import NUS_RX_UUID, NUS_TX_UUID
 from ..ble.pybricks import (
+    DEVICE_NAME_UUID,
     FW_REV_UUID,
     PNP_ID_UUID,
     PYBRICKS_COMMAND_EVENT_UUID,
@@ -142,6 +149,140 @@ def handler(_, data):
         return await self._client.start_notify(PYBRICKS_COMMAND_EVENT_UUID, handler)
 
 
+class _USBTransport(_Transport):
+    _device: USBDevice
+    _disconnected_callback: Callable
+    _ep_in: Endpoint
+    _ep_out: Endpoint
+    _notify_callbacks = {}
+    _monitor_task: asyncio.Task
+
+    def __init__(self, device: USBDevice):
+        self._device = device
+
+    async def connect(self, disconnected_callback: Callable) -> None:
+        self._disconnected_callback = disconnected_callback
+        self._device.set_configuration()
+
+        # Save input and output endpoints
+        cfg = self._device.get_active_configuration()
+        intf = cfg[(0, 0)]
+        self._ep_in = find_descriptor(
+            intf,
+            custom_match=lambda e: endpoint_direction(e.bEndpointAddress)
+            == ENDPOINT_IN,
+        )
+        self._ep_out = find_descriptor(
+            intf,
+            custom_match=lambda e: endpoint_direction(e.bEndpointAddress)
+            == ENDPOINT_OUT,
+        )
+
+        # Set write size to endpoint packet size minus length of UUID
+        self._max_write_size = self._ep_out.wMaxPacketSize - 16
+
+        # Get length of BOS descriptor
+        bos_descriptor = get_descriptor(self._device, 5, 0x0F, 0)
+        (ofst, _, bos_len, _) = struct.unpack("<BBHB", bos_descriptor)
+
+        # Get full BOS descriptor
+        bos_descriptor = get_descriptor(self._device, bos_len, 0x0F, 0)
+
+        while ofst < bos_len:
+            (len, desc_type, cap_type) = struct.unpack_from(
+                "<BBB", bos_descriptor, offset=ofst
+            )
+
+            if desc_type != 0x10:
+                logger.error("Expected Device Capability descriptor")
+                exit(1)
+
+            # Look for platform descriptors
+            if cap_type == 0x05:
+                uuid_bytes = bos_descriptor[ofst + 4 : ofst + 4 + 16]
+                uuid_str = str(UUID(bytes_le=bytes(uuid_bytes)))
+
+                if uuid_str == DEVICE_NAME_UUID:
+                    self._device_name = bytearray(
+                        bos_descriptor[ofst + 20 : ofst + len]
+                    ).decode()
+                    print("Connected to hub '" + self._device_name + "'")
+
+                elif uuid_str == FW_REV_UUID:
+                    fw_version = bytearray(bos_descriptor[ofst + 20 : ofst + len])
+                    self._fw_version = Version(fw_version.decode())
+
+                elif uuid_str == SW_REV_UUID:
+                    protocol_version = bytearray(bos_descriptor[ofst + 20 : ofst + len])
+                    self._protocol_version = semver.VersionInfo.parse(
+                        protocol_version.decode()
+                    )
+
+                elif uuid_str == PYBRICKS_HUB_CAPABILITIES_UUID:
+                    caps = bytearray(bos_descriptor[ofst + 20 : ofst + len])
+                    (
+                        _,
+                        self._capability_flags,
+                        self._max_user_program_size,
+                    ) = unpack_hub_capabilities(caps)
+
+            ofst += len
+
+        self._monitor_task = asyncio.create_task(self._monitor_usb())
+
+    async def disconnect(self) -> None:
+        self._monitor_task.cancel()
+        self._disconnected_callback()
+
+    async def get_firmware_version(self) -> Version:
+        return self._fw_version
+
+    async def get_protocol_version(self) -> Version:
+        return self._protocol_version
+
+    async def get_hub_type(self) -> Tuple[HubKind, int]:
+        return (HubKind.TECHNIC_LARGE, 0)
+
+    async def get_hub_capabilities(self) -> Tuple[int, HubCapabilityFlag, int]:
+        return (
+            self._max_write_size,
+            self._capability_flags,
+            self._max_user_program_size,
+        )
+
+    async def send_command(self, command: bytes) -> None:
+        self._ep_out.write(UUID(PYBRICKS_COMMAND_EVENT_UUID).bytes_le + command)
+
+    async def set_nus_handler(self, callback: Callable) -> None:
+        pass
+
+    async def set_service_handler(self, callback: Callable) -> None:
+        self._notify_callbacks[PYBRICKS_COMMAND_EVENT_UUID] = callback
+
+    async def _monitor_usb(self):
+        loop = asyncio.get_running_loop()
+
+        while True:
+            msg = await loop.run_in_executor(None, self._read_usb)
+
+            if msg is None:
+                continue
+
+            if len(msg) > 16:
+                uuid = str(UUID(bytes_le=bytes(msg[0:16])))
+                if uuid in self._notify_callbacks:
+                    callback = self._notify_callbacks[uuid]
+                    if callback:
+                        callback(bytes(msg[16:]))
+
+    def _read_usb(self):
+        try:
+            msg = self._ep_in.read(self._ep_in.wMaxPacketSize)
+            return msg
+        except USBTimeoutError:
+            return None
+
+
 class PybricksHub:
     EOL = b"\r\n"  # MicroPython EOL
 
@@ -178,7 +319,7 @@ class PybricksHub:
 
     _transport: _Transport
 
-    def __init__(self, device: BLEDevice):
+    def __init__(self, device):
         self.connection_state_observable = BehaviorSubject(ConnectionState.DISCONNECTED)
         self.status_observable = BehaviorSubject(StatusFlag(0))
         self._stdout_subject = Subject()
@@ -220,7 +361,10 @@ def __init__(self, device: BLEDevice):
         # File handle for logging
         self.log_file = None
 
-        self._transport = _BLETransport(device)
+        if isinstance(device, BLEDevice):
+            self._transport = _BLETransport(device)
+        elif isinstance(device, USBDevice):
+            self._transport = _USBTransport(device)
 
     @property
     def stdout_observable(self) -> Observable[bytes]: