cli.lwp.repl: new command line tool for LWP3

This adds an interactive REPL tool for sending and receiving messages
using the LWP3 protocol.

Author: dlech

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Added
 - `pybricksdev.ble.lwp3.bytecodes` module.
 - `pybricksdev.ble.lwp3.messages` module.
+- `pybricksdev lwp3 repl` command line tool.
 
 ## [1.0.0-alpha.8] - 2021-05-18
 ## Added

poetry.lock

@@ -258,6 +258,43 @@ def run(self, args: argparse.Namespace):
         return self.subparsers.choices[args.action].tool.run(args)
 
 
+class LWP3Repl(Tool):
+    def add_parser(self, subparsers: argparse._SubParsersAction):
+        parser = subparsers.add_parser(
+            "repl",
+            help="interactive REPL for sending and receiving LWP3 messages",
+        )
+        parser.tool = self
+
+    def run(self, args: argparse.Namespace):
+        from .lwp3.repl import setup_repl_logging, repl
+
+        setup_repl_logging()
+        return repl()
+
+
+class LWP3(Tool):
+    def add_parser(self, subparsers: argparse._SubParsersAction):
+        self.parser = subparsers.add_parser(
+            "lwp3", help="interact with devices using LWP3"
+        )
+        self.parser.tool = self
+        self.subparsers = self.parser.add_subparsers(
+            metavar="<lwp3-tool>", dest="lwp3_tool", help="the tool to run"
+        )
+
+        for tool in (LWP3Repl(),):
+            tool.add_parser(self.subparsers)
+
+    def run(self, args: argparse.Namespace):
+        if args.lwp3_tool not in self.subparsers.choices:
+            self.parser.error(
+                f'Missing name of tool: {"|".join(self.subparsers.choices.keys())}'
+            )
+
+        return self.subparsers.choices[args.lwp3_tool].tool.run(args)
+
+
 class Udev(Tool):
     def add_parser(self, subparsers: argparse._SubParsersAction):
         parser = subparsers.add_parser("udev", help="print udev rules to stdout")
@@ -293,7 +330,7 @@ def main():
         help="the tool to use",
     )
 
-    for tool in Compile(), Run(), Flash(), DFU(), Udev():
+    for tool in Compile(), Run(), Flash(), DFU(), LWP3(), Udev():
         tool.add_parser(subparsers)
 
     argcomplete.autocomplete(parser)

pybricksdev/cli/__init__.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2021 The Pybricks Authors
+
+"""
+The :mod:`pybricks.cli.lwp3.repl` module provides a command line interface
+for connecting to a device and sending and receiving LWP3 messages.
+"""
+
+import asyncio
+import inspect
+import logging
+import os
+import struct
+from pathlib import Path
+
+from appdirs import user_cache_dir
+from bleak import BleakClient, BleakScanner
+from bleak.backends.device import BLEDevice
+from bleak.backends.scanner import AdvertisementData
+from prompt_toolkit import PromptSession
+from prompt_toolkit.history import FileHistory
+from prompt_toolkit.patch_stdout import StdoutProxy, patch_stdout
+
+from ...ble.lwp3 import (
+    LEGO_CID,
+    LWP3_HUB_CHARACTERISTIC_UUID,
+    LWP3_HUB_SERVICE_UUID,
+    bytecodes,
+    messages,
+)
+from ...ble.lwp3.bytecodes import Capabilities, HubKind, LastNetwork, Status
+from ...ble.lwp3.messages import AbstractMessage, parse_message
+
+logger = logging.getLogger(__name__)
+history_file = Path(user_cache_dir("pybricksdev"), "lwp3-explorer-history.txt")
+
+# Get names that are valid for evaluating on the REPL.
+# This hides built-in functions to avoid arbitrary code execution.
+_eval_pool = {"__builtins__": {}}
+
+# TODO: these dicts can be used for tab completion on the REPL
+
+# The first groups is any type from bytecodes that inherits from int (includes
+# enums/flags) or bytes.
+_eval_pool.update(
+    {
+        k: v
+        for k, v in bytecodes.__dict__.items()
+        if inspect.isclass(v)
+        and v.__module__ == bytecodes.__name__
+        and (issubclass(v, int) or issubclass(v, bytes))
+    }
+)
+
+# The second group are all of the non-abstract message types from the messages module.
+_eval_pool.update(
+    {
+        k: v
+        for k, v in messages.__dict__.items()
+        if inspect.isclass(v)
+        and issubclass(v, AbstractMessage)
+        and not inspect.isabstract(v)
+    }
+)
+
+
+async def repl() -> None:
+    """
+    Provides an interactive REPL for sending and receiving LWP3 messages.
+    """
+    os.makedirs(history_file.parent, exist_ok=True)
+    session = PromptSession(history=FileHistory(history_file))
+
+    queue = asyncio.Queue()
+
+    def callback(dev: BLEDevice, adv: AdvertisementData) -> None:
+        if LWP3_HUB_SERVICE_UUID.lower() not in adv.service_uuids:
+            return
+
+        mfg_data = adv.manufacturer_data[LEGO_CID]
+        button, kind, cap, last_net, status, opt = struct.unpack("<6B", mfg_data)
+        button = bool(button)
+        kind = HubKind(kind)
+        cap = Capabilities(cap)
+        last_net = LastNetwork(last_net)
+        status = Status(status)
+        logger.debug(
+            "button: %s, kind: %s, cap: %s, last net: %s, status: %s, option: %s",
+            button,
+            kind,
+            cap,
+            last_net,
+            status,
+            opt,
+        )
+
+        if not (cap & Capabilities.REMOTE):
+            return
+
+        queue.put_nowait(dev)
+
+    async with BleakScanner(detection_callback=callback):
+        logger.info("scanning...")
+        device = await queue.get()
+        logger.info("found device")
+
+    def handle_disconnect(client: BleakClient):
+        logger.info("disconnected")
+
+    async with BleakClient(device, disconnected_callback=handle_disconnect) as client:
+        logger.info("connected")
+
+        def handle_notify(handle, value):
+            try:
+                msg = parse_message(value)
+            except Exception as ex:
+                logger.warning("failed to parse message: %s", value, exc_info=ex)
+            else:
+                logger.info("received: %s", msg)
+
+        await client.start_notify(LWP3_HUB_CHARACTERISTIC_UUID, handle_notify)
+
+        # welcome is delayed to allow initial log messages to settle.
+        async def welcome():
+            await asyncio.sleep(1)
+            print("Type message and press ENTER to send. Press CTRL+D to exit.")
+
+        asyncio.ensure_future(welcome())
+
+        while True:
+            with patch_stdout():
+                try:
+                    result = await session.prompt_async(">>> ")
+                except KeyboardInterrupt:
+                    # CTRL+C ignores the line
+                    continue
+                except EOFError:
+                    # CTRL+D exits the program
+                    break
+            try:
+                msg = eval(result, _eval_pool)
+                if not isinstance(msg, AbstractMessage):
+                    raise ValueError("not a message object")
+            except Exception:
+                logger.exception("bad input:")
+            else:
+                logger.info("sending: %s", msg)
+                await client.write_gatt_char(LWP3_HUB_CHARACTERISTIC_UUID, msg)
+
+        logger.info("disconnecting...")
+
+
+def setup_repl_logging() -> None:
+    """
+    Overrides logging as needed for :func:`repl`.
+    """
+    logging.basicConfig(
+        stream=StdoutProxy(),
+        format="[%(asctime)s.%(msecs)03d] %(levelname)s: %(message)s",
+        datefmt="%H:%M:%S",
+        force=True,
+    )
+    logger.setLevel(logging.INFO)
+
+
+if __name__ == "__main__":
+    setup_repl_logging()
+    asyncio.run(repl())

pybricksdev/cli/lwp3/__init__.py

@@ -37,6 +37,8 @@ tqdm = "^4.46.1"
 validators = "^0.18.2"
 pyusb = "^1.0.2"
 semver = "^2.13.0"
+appdirs = "^1.4.4"
+prompt-toolkit = "^3.0.18"
 
 [tool.poetry.dev-dependencies]
 black = {version = "^21.5b1", allow-prereleases = true}