cli: modify flash command to reboot in update mode

This adds a feature the the CLI `flash` command to scan for hubs not
in bootloader mode running official LEGO firmware or Pybricks firmware.
It will reboot these hubs in bootloader mode and proceed with flashing
the firmware.

Author: dlech

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Moved ``Bootloader*`` to new ``ble.lwp3.bootloader`` module.
+- ``pybricksdev flash`` will now discover hubs running official LEGO firmware
+  or Pybricks firmware and reboot in bootloader mode automatically.
 
 ## [1.0.0-alpha.21] - 2022-01-12
 

pybricksdev/cli/flash.py

@@ -3,15 +3,285 @@
 
 import asyncio
 import json
+import logging
+import sys
 import zipfile
+from tempfile import NamedTemporaryFile
 from typing import BinaryIO, Optional
 
-from ..ble import find_device
-from ..ble.lwp3 import LWP3_BOOTLOADER_SERVICE_UUID
-from ..ble.lwp3.bytecodes import HubKind
+from bleak import BleakClient, BleakScanner
+from bleak.backends.device import BLEDevice
+from bleak.backends.scanner import AdvertisementData
+
+from ..ble.lwp3 import (
+    LEGO_CID,
+    LWP3_BOOTLOADER_SERVICE_UUID,
+    LWP3_HUB_CHARACTERISTIC_UUID,
+    LWP3_HUB_SERVICE_UUID,
+)
+from ..ble.lwp3 import AdvertisementData as HubAdvertisementData
+from ..ble.lwp3.bootloader import BootloaderAdvertisementData
+from ..ble.lwp3.bytecodes import HubKind, HubProperty
+from ..ble.lwp3.messages import (
+    FirmwareUpdateMessage,
+    HubPropertyRequestUpdate,
+    HubPropertyUpdate,
+    parse_message,
+)
+from ..ble.nus import NUS_RX_UUID, NUS_TX_UUID
+from ..ble.pybricks import (
+    FW_REV_UUID,
+    PNP_ID_UUID,
+    PYBRICKS_SERVICE_UUID,
+    unpack_pnp_id,
+)
+from ..compile import compile_file
 from ..connections import REPLHub
 from ..dfu import flash_dfu
 from ..flash import BootloaderConnection, create_firmware
+from ..tools import chunk
+from ..tools.checksum import xor_bytes
+
+logger = logging.getLogger(__name__)
+
+
+REBOOT_SCRIPT = """
+from pybricks.hubs import ThisHub
+from pybricks.tools import wait
+
+hub = ThisHub()
+
+# without delay, hub will reboot before we receive last checksum
+wait(500)
+hub.system.reset(2)
+"""
+
+
+def match_hub(
+    hub_kind: HubKind, hub_name: Optional[str], adv: AdvertisementData
+) -> bool:
+    """
+    Advertisement data matching function for filtering supported hubs.
+
+    Args:
+        hub_kind: The hub type ID to match.
+        hub_name: An optional name to match to the "local name" in the
+            advertisement data.
+        adv: The advertisemet data to check.
+
+    Returns:
+        ``True`` if *adv* matches the criteria, otherwise ``False``.
+    """
+    # hub name filtering is optional
+    if hub_name and adv.local_name != hub_name:
+        return False
+
+    # LEGO firmware uses manufacturer-specific data
+
+    lego_data = adv.manufacturer_data.get(LEGO_CID)
+
+    if lego_data:
+        if LWP3_BOOTLOADER_SERVICE_UUID in adv.service_uuids:
+            bl_data = BootloaderAdvertisementData(lego_data)
+            return bl_data.hub_kind == hub_kind
+
+        if LWP3_HUB_SERVICE_UUID in adv.service_uuids:
+            hub_data = HubAdvertisementData(lego_data)
+            return hub_data.hub_kind == hub_kind
+
+    # Pybricks firmware uses Device Information service data
+
+    pnp_id_data = adv.service_data.get(PNP_ID_UUID)
+
+    if pnp_id_data and PYBRICKS_SERVICE_UUID in adv.service_uuids:
+        _, _, pid, _ = unpack_pnp_id(pnp_id_data)
+        return pid == hub_kind
+
+    return False
+
+
+async def download_and_run(client: BleakClient, script: str) -> None:
+    """
+    Downloads and runs a script on a hub running Pybricks firmware.
+
+    Args:
+        client: The Bluetooth connection to the hub.
+        script: The script to be compiled and run.
+    """
+    with NamedTemporaryFile("w", suffix=".py") as temp:
+        temp.write(script)
+
+        # file has to be closed so mpy-cross can open it
+        temp.file.close()
+
+        mpy = await compile_file(temp.name)
+
+    recv_queue = asyncio.Queue()
+
+    def on_notify(_h, data: bytes):
+        recv_queue.put_nowait(data)
+
+    # BOOST Move hub has hardware limit of MTU == 23 so it has to have data
+    # split into smaller chunks
+    write_size = 20 if client.mtu_size < 100 else 100
+
+    async def write_chunk(data: bytes):
+        """
+        Writes a chunk of data and waits for a checksum reply.
+
+        Args:
+            data: The data.
+
+        Raises:
+            RuntimeError: If the returned checksum did not match.
+            asyncio.TimeoutError: If no reply was received.
+        """
+        checksum = xor_bytes(data, 0)
+
+        for c in chunk(data, write_size):
+            await client.write_gatt_char(NUS_RX_UUID, c)
+
+        reply: bytes = await asyncio.wait_for(recv_queue.get(), 1)
+
+        if reply[0] != checksum:
+            raise RuntimeError("bad checksum")
+
+    await client.start_notify(NUS_TX_UUID, on_notify)
+
+    # communication protocol is write file size, then send file in 100 byte chunks
+    try:
+        await write_chunk(len(mpy).to_bytes(4, "little"))
+
+        for c in chunk(mpy, 100):
+            await write_chunk(c)
+
+    finally:
+        await client.stop_notify(NUS_TX_UUID)
+
+
+async def reboot_official_to_bootloader(hub_kind: HubKind, device: BLEDevice) -> None:
+    """
+    Connects to a hub running official LEGO firmware and sends a message to
+    reboot in firmware update mode.
+    """
+    async with BleakClient(device) as client:
+
+        # give bluetooth stack time to settle
+        await asyncio.sleep(1)
+
+        fw_ver_future = asyncio.get_running_loop().create_future()
+
+        def on_notify(_h, data: bytes):
+            msg = parse_message(data)
+
+            logger.debug("%s", str(msg))
+
+            if (
+                isinstance(msg, HubPropertyUpdate)
+                and msg.prop == HubProperty.FW_VERSION
+            ):
+                fw_ver_future.set_result(msg.value)
+
+        await client.start_notify(LWP3_HUB_CHARACTERISTIC_UUID, on_notify)
+        await client.write_gatt_char(
+            LWP3_HUB_CHARACTERISTIC_UUID,
+            HubPropertyRequestUpdate(HubProperty.FW_VERSION),
+            # work around city hub bluetooth bug on linux
+            response=hub_kind == HubKind.CITY,
+        )
+
+        fw_ver = await asyncio.wait_for(fw_ver_future, 5)
+        print(f"Hub is running firmware v{fw_ver}.")
+
+        print("Rebooting in update mode...")
+
+        await client.write_gatt_char(
+            LWP3_HUB_CHARACTERISTIC_UUID, FirmwareUpdateMessage()
+        )
+
+
+async def reboot_pybricks_to_bootloader(hub_kind: HubKind, device: BLEDevice) -> None:
+    """
+    Connects to a hub running Pybricks firmware and sends a message to
+    reboot in firmware update mode.
+    """
+    async with BleakClient(device) as client:
+        # Work around BlueZ limitation.
+        if client.__class__.__name__ == "BleakClientBlueZDBus":
+            client._mtu_size = 23 if hub_kind == HubKind.BOOST else 158
+
+        # give bluetooth stack time to settle
+        await asyncio.sleep(1)
+
+        fw_ver = await client.read_gatt_char(FW_REV_UUID)
+        fw_ver = fw_ver.decode()
+        print(f"Hub is running firmware v{fw_ver}.")
+
+        print("Rebooting in update mode...")
+
+        await download_and_run(client, REBOOT_SCRIPT)
+
+
+async def flash_ble(
+    hub_kind: HubKind, hub_name: Optional[str], firmware: bytes, metadata: dict
+):
+    """
+    Flashes firmware to the hub using Bluetooth Low Energy.
+
+    The hub has to be advertising and can be running official LEGO firmware,
+    Pybricks firmware or be in bootloader mode.
+
+    Args:
+        hub_kind: The hub type ID. Only hubs matching this ID will be discovered.
+        hub_name: If given, only hubs that advertise this name will be disovered.
+        firmware: The raw firmware binary blob.
+        metadata: The firmware metadata from the firmware.zip file.
+    """
+
+    print(f"Searching for {repr(hub_name) if hub_name else hub_kind.name} hub...")
+
+    # scan for hubs in bootloader mode, running official LEGO firmware or
+    # running Pybricks firmware
+
+    device = await BleakScanner.find_device_by_filter(
+        lambda _d, a: match_hub(hub_kind, hub_name, a),
+        service_uuids=[
+            LWP3_BOOTLOADER_SERVICE_UUID,
+            LWP3_HUB_SERVICE_UUID,
+            PYBRICKS_SERVICE_UUID,
+        ],
+    )
+
+    if device is None:
+        print("timed out", file=sys.stderr)
+        return
+
+    # if not already in bootlaoder mode, we need to reboot into bootloader mode
+    if LWP3_HUB_SERVICE_UUID in device.metadata["uuids"]:
+        print("Found hub running official LEGO firmare.")
+        await reboot_official_to_bootloader(hub_kind, device)
+    elif PYBRICKS_SERVICE_UUID in device.metadata["uuids"]:
+        print("Found hub running Pybricks firmare.")
+        await reboot_pybricks_to_bootloader(hub_kind, device)
+
+    # if not previously in bootlaoder mode, scan again, this time only for bootloader
+    if LWP3_BOOTLOADER_SERVICE_UUID not in device.metadata["uuids"]:
+        device = await BleakScanner.find_device_by_filter(
+            lambda _d, a: match_hub(hub_kind, hub_name, a),
+            service_uuids=[
+                LWP3_BOOTLOADER_SERVICE_UUID,
+            ],
+        )
+
+        if device is None:
+            print("timed out", file=sys.stderr)
+            return
+
+    print("Found:", device)
+    updater = BootloaderConnection()
+    await updater.connect(device)
+    print("Erasing flash and starting update")
+    await updater.flash(firmware, metadata)
 
 
 async def flash_firmware(firmware_zip: BinaryIO, hub_name: Optional[str]) -> None:
@@ -26,8 +296,9 @@ async def flash_firmware(firmware_zip: BinaryIO, hub_name: Optional[str]) -> Non
     print("Creating firmware...")
 
     firmware, metadata = await create_firmware(firmware_zip, hub_name)
+    hub_kind = HubKind(metadata["device-id"])
 
-    if metadata["device-id"] in (HubKind.TECHNIC_SMALL, HubKind.TECHNIC_LARGE):
+    if hub_kind in (HubKind.TECHNIC_SMALL, HubKind.TECHNIC_LARGE):
         try:
             # Connect to the hub and exit the runtime.
             hub = REPLHub()
@@ -60,16 +331,4 @@ async def flash_firmware(firmware_zip: BinaryIO, hub_name: Optional[str]) -> Non
             print("Could not find hub in standard firmware mode. Trying DFU.")
             flash_dfu(firmware, metadata)
     else:
-        print("Searching for LEGO Bootloader...")
-
-        try:
-            device = await find_device(service=LWP3_BOOTLOADER_SERVICE_UUID)
-        except asyncio.TimeoutError:
-            print("timed out")
-            return
-
-        print("Found:", device)
-        updater = BootloaderConnection()
-        await updater.connect(device)
-        print("Erasing flash and starting update")
-        await updater.flash(firmware, metadata)
+        await flash_ble(hub_kind, hub_name, firmware, metadata)