foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/robot/ble_enumerate.py

@@ -0,0 +1,506 @@
+#!/usr/bin/env python3
+"""
+Comprehensive BLE Enumeration for Robots.
+
+This module provides systematic BLE device enumeration that can discover:
+1. All GATT services and characteristics
+2. Which characteristics accept writes/reads/notifications
+3. What command protocols the device responds to
+4. Visual confirmation of command execution (if camera available)
+
+Designed to work with any BLE-enabled robot, not just MechDog.
+"""
+
+import asyncio
+import json
+import struct
+import time
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Optional, Callable, Any
+
+from bleak import BleakClient, BleakScanner
+from bleak.backends.characteristic import BleakGATTCharacteristic
+
+
+class CommandResult(Enum):
+    """Result of sending a command."""
+    NO_RESPONSE = "no_response"
+    RESPONSE = "response"
+    ERROR = "error"
+    EXECUTED = "executed"  # No response but visual/sensor confirmation
+
+
+@dataclass
+class CharacteristicInfo:
+    """Information about a BLE characteristic."""
+    uuid: str
+    handle: int
+    properties: list[str]
+    descriptors: list[str] = field(default_factory=list)
+    accepts_write: bool = False
+    accepts_read: bool = False
+    supports_notify: bool = False
+    read_value: Optional[bytes] = None
+    write_test_result: Optional[str] = None
+
+
+@dataclass
+class ServiceInfo:
+    """Information about a BLE service."""
+    uuid: str
+    handle: int
+    characteristics: list[CharacteristicInfo] = field(default_factory=list)
+
+
+@dataclass
+class CommandProbe:
+    """A command to probe and its result."""
+    command: bytes
+    description: str
+    response: Optional[bytes] = None
+    response_text: Optional[str] = None
+    result: CommandResult = CommandResult.NO_RESPONSE
+    latency_ms: float = 0.0
+
+
+@dataclass
+class ProtocolProbe:
+    """Results of probing a specific protocol."""
+    name: str
+    description: str
+    commands: list[CommandProbe] = field(default_factory=list)
+    working: bool = False
+
+
+@dataclass
+class BLEEnumerationReport:
+    """Complete enumeration report for a BLE device."""
+    device_name: str
+    device_address: str
+    timestamp: str
+    services: list[ServiceInfo] = field(default_factory=list)
+    protocols: list[ProtocolProbe] = field(default_factory=list)
+    primary_write_char: Optional[str] = None
+    primary_notify_char: Optional[str] = None
+    working_commands: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for JSON serialization."""
+        data = asdict(self)
+
+        # Convert enum values, bytes, and bytearray to strings
+        def convert_types(obj):
+            if isinstance(obj, Enum):
+                return obj.value
+            if isinstance(obj, (bytes, bytearray)):
+                return obj.hex()
+            if isinstance(obj, dict):
+                return {k: convert_types(v) for k, v in obj.items()}
+            if isinstance(obj, list):
+                return [convert_types(i) for i in obj]
+            return obj
+
+        return convert_types(data)
+
+    def to_json(self, path: Path) -> None:
+        """Save report to JSON file."""
+        # Convert bytes to hex strings for JSON serialization
+        data = self.to_dict()
+
+        def convert_bytes(obj):
+            if isinstance(obj, bytes):
+                return obj.hex()
+            if isinstance(obj, dict):
+                return {k: convert_bytes(v) for k, v in obj.items()}
+            if isinstance(obj, list):
+                return [convert_bytes(i) for i in obj]
+            return obj
+
+        data = convert_bytes(data)
+        path.write_text(json.dumps(data, indent=2))
+
+
+class BLEEnumerator:
+    """Systematic BLE device enumeration."""
+
+    # Common robot BLE protocols to probe
+    PROTOCOLS = {
+        "cmd_pipe": {
+            "description": "CMD|X|$ pipe-delimited protocol (HiWonder/Hiwonder robots)",
+            "commands": [
+                (b"CMD|6|$", "Battery query"),
+                (b"CMD|0|0|$", "Mode 0 (stop)"),
+                (b"CMD|1|0|$", "Pose mode"),
+                (b"CMD|2|1|$", "Action 1"),
+                (b"CMD|3|1|$", "Gait mode 1"),
+                (b"CMD|4|0|0|$", "Direction stop"),
+                (b"CMD|5|0|0|0|$", "Attitude neutral"),
+                (b"CMD|7|6|$", "Arm command 6"),
+                (b"CMD|7|7|$", "Arm command 7"),
+            ]
+        },
+        "micropython_repl": {
+            "description": "MicroPython REPL protocol",
+            "commands": [
+                (b"\x03", "Ctrl+C interrupt"),
+                (b"\x02", "Ctrl+B friendly REPL"),
+                (b"\r\n", "Empty line (prompt)"),
+                (b"print('test')\r\n", "Python print"),
+                (b"2+2\r\n", "Python eval"),
+            ]
+        },
+        "hiwonder_servo": {
+            "description": "HiWonder servo binary protocol (0x55 0x55 header)",
+            "commands": [
+                # Servo 1 read position
+                (bytes([0x55, 0x55, 0x01, 0x03, 0x1C, 0xE4]), "Servo 1 read pos"),
+                # Servo 254 (broadcast) read
+                (bytes([0x55, 0x55, 0xFE, 0x03, 0x1C, 0xE7]), "Broadcast read"),
+            ]
+        },
+        "json": {
+            "description": "JSON command protocol",
+            "commands": [
+                (b'{"cmd":"status"}', "JSON status"),
+                (b'{"cmd":"battery"}', "JSON battery"),
+                (b'{"action":"stop"}', "JSON stop"),
+            ]
+        },
+        "at_commands": {
+            "description": "AT command protocol",
+            "commands": [
+                (b"AT\r\n", "AT ping"),
+                (b"AT+VERSION\r\n", "AT version"),
+                (b"AT+NAME\r\n", "AT name"),
+            ]
+        },
+    }
+
+    def __init__(
+        self,
+        device_address: str,
+        timeout: float = 15.0,
+        command_delay: float = 0.3,
+        verbose: bool = True,
+        visual_callback: Optional[Callable[[], float]] = None,
+    ):
+        """
+        Initialize BLE enumerator.
+
+        Args:
+            device_address: BLE MAC address or UUID
+            timeout: Connection timeout in seconds
+            command_delay: Delay between commands in seconds
+            verbose: Print progress information
+            visual_callback: Optional callback that returns visual change percentage
+        """
+        self.device_address = device_address
+        self.timeout = timeout
+        self.command_delay = command_delay
+        self.verbose = verbose
+        self.visual_callback = visual_callback
+        self.report = None
+
+    def log(self, msg: str) -> None:
+        """Print if verbose mode enabled."""
+        if self.verbose:
+            print(msg)
+
+    async def scan_for_device(self) -> Optional[Any]:
+        """Scan for the target device."""
+        self.log(f"Scanning for device: {self.device_address}")
+        devices = await BleakScanner.discover(timeout=5.0)
+
+        for d in devices:
+            if d.address == self.device_address or d.name == self.device_address:
+                self.log(f"  Found: {d.name} ({d.address})")
+                return d
+
+        # Try case-insensitive name match
+        for d in devices:
+            if d.name and self.device_address.lower() in d.name.lower():
+                self.log(f"  Found by name: {d.name} ({d.address})")
+                return d
+
+        self.log("  Device not found")
+        return None
+
+    async def enumerate_services(self, client: BleakClient) -> list[ServiceInfo]:
+        """Enumerate all GATT services and characteristics."""
+        services = []
+
+        for service in client.services:
+            svc_info = ServiceInfo(
+                uuid=str(service.uuid),
+                handle=service.handle,
+            )
+
+            for char in service.characteristics:
+                char_info = CharacteristicInfo(
+                    uuid=str(char.uuid),
+                    handle=char.handle,
+                    properties=char.properties,
+                    descriptors=[str(d.uuid) for d in char.descriptors],
+                    accepts_write="write" in char.properties or "write-without-response" in char.properties,
+                    accepts_read="read" in char.properties,
+                    supports_notify="notify" in char.properties or "indicate" in char.properties,
+                )
+
+                # Try reading if readable
+                if char_info.accepts_read:
+                    try:
+                        value = await client.read_gatt_char(char.uuid)
+                        char_info.read_value = value
+                    except Exception:
+                        pass
+
+                svc_info.characteristics.append(char_info)
+
+            services.append(svc_info)
+
+        return services
+
+    async def probe_protocol(
+        self,
+        client: BleakClient,
+        write_char: str,
+        notify_char: str,
+        protocol_name: str,
+        protocol_info: dict,
+    ) -> ProtocolProbe:
+        """Probe a specific protocol."""
+        probe = ProtocolProbe(
+            name=protocol_name,
+            description=protocol_info["description"],
+        )
+
+        responses = []
+
+        def handler(sender, data):
+            responses.append((time.time(), data))
+            if self.verbose:
+                try:
+                    self.log(f"      << {data.decode('ascii', errors='replace')}")
+                except:
+                    self.log(f"      << {data.hex()}")
+
+        # Start notifications
+        try:
+            await client.start_notify(notify_char, handler)
+            await asyncio.sleep(0.5)  # Give time to establish
+        except Exception as e:
+            self.log(f"    Failed to start notify: {e}")
+
+        await asyncio.sleep(0.2)
+
+        for cmd_bytes, cmd_desc in protocol_info["commands"]:
+            cmd_probe = CommandProbe(
+                command=cmd_bytes,
+                description=cmd_desc,
+            )
+
+            responses.clear()
+            start = time.time()
+
+            try:
+                await client.write_gatt_char(write_char, cmd_bytes)
+                await asyncio.sleep(self.command_delay)
+
+                if responses:
+                    latency = (responses[0][0] - start) * 1000
+                    resp_data = responses[0][1]
+
+                    cmd_probe.response = resp_data
+                    cmd_probe.latency_ms = latency
+                    cmd_probe.result = CommandResult.RESPONSE
+
+                    try:
+                        cmd_probe.response_text = resp_data.decode('ascii', errors='replace')
+                    except:
+                        cmd_probe.response_text = resp_data.hex()
+
+                    probe.working = True
+                else:
+                    # Check for visual movement if callback provided
+                    if self.visual_callback:
+                        change = self.visual_callback()
+                        if change > 3.0:  # 3% threshold
+                            cmd_probe.result = CommandResult.EXECUTED
+                            probe.working = True
+                    else:
+                        cmd_probe.result = CommandResult.NO_RESPONSE
+
+            except Exception as e:
+                cmd_probe.result = CommandResult.ERROR
+                cmd_probe.response_text = str(e)
+
+            probe.commands.append(cmd_probe)
+
+        try:
+            await client.stop_notify(notify_char)
+        except Exception:
+            pass
+
+        return probe
+
+    async def enumerate(self) -> BLEEnumerationReport:
+        """Run full BLE enumeration."""
+        device = await self.scan_for_device()
+        if not device:
+            raise RuntimeError(f"Device {self.device_address} not found")
+
+        self.report = BLEEnumerationReport(
+            device_name=device.name or "Unknown",
+            device_address=device.address,
+            timestamp=datetime.now().isoformat(),
+        )
+
+        self.log(f"\nConnecting to {device.name}...")
+
+        async with BleakClient(device, timeout=self.timeout) as client:
+            # Phase 1: Enumerate services
+            self.log("\n[1] Enumerating GATT services...")
+            self.report.services = await self.enumerate_services(client)
+
+            for svc in self.report.services:
+                self.log(f"  Service: {svc.uuid}")
+                for char in svc.characteristics:
+                    props = ", ".join(char.properties)
+                    self.log(f"    Char: {char.uuid} [{props}]")
+                    if char.read_value:
+                        try:
+                            self.log(f"      Value: {char.read_value.decode('ascii', errors='replace')}")
+                        except:
+                            self.log(f"      Value: {char.read_value.hex()}")
+
+            # Find write and notify characteristics
+            write_chars = []
+            notify_chars = []
+
+            for svc in self.report.services:
+                for char in svc.characteristics:
+                    if char.accepts_write:
+                        write_chars.append(char.uuid)
+                    if char.supports_notify:
+                        notify_chars.append(char.uuid)
+
+            if not write_chars:
+                self.report.notes.append("No writable characteristics found")
+                return self.report
+
+            # Use first write/notify pair (common pattern)
+            write_char = write_chars[0]
+            notify_char = notify_chars[0] if notify_chars else write_chars[0]
+
+            self.report.primary_write_char = write_char
+            self.report.primary_notify_char = notify_char
+
+            self.log(f"\n  Primary write: {write_char}")
+            self.log(f"  Primary notify: {notify_char}")
+
+            # Phase 2: Probe protocols
+            self.log("\n[2] Probing command protocols...")
+
+            for proto_name, proto_info in self.PROTOCOLS.items():
+                self.log(f"\n  Testing {proto_name}...")
+                probe = await self.probe_protocol(
+                    client, write_char, notify_char, proto_name, proto_info
+                )
+                self.report.protocols.append(probe)
+
+                if probe.working:
+                    self.log(f"    ✓ {proto_name} WORKING")
+                    for cmd in probe.commands:
+                        if cmd.result in (CommandResult.RESPONSE, CommandResult.EXECUTED):
+                            self.log(f"      {cmd.description}: {cmd.response_text or 'executed'}")
+                            self.report.working_commands.append(cmd.description)
+                else:
+                    self.log(f"    ✗ {proto_name} not responding")
+
+            # Phase 3: Deep scan of working protocols
+            if any(p.working for p in self.report.protocols):
+                self.log("\n[3] Deep scanning working protocols...")
+
+                # If CMD protocol works, scan more functions
+                cmd_proto = next((p for p in self.report.protocols if p.name == "cmd_pipe" and p.working), None)
+                if cmd_proto:
+                    self.log("  Scanning CMD|X|$ functions 0-20...")
+
+                    responses = []
+                    def handler(s, d):
+                        responses.append(d)
+
+                    await client.start_notify(notify_char, handler)
+
+                    for func in range(21):
+                        responses.clear()
+                        cmd = f"CMD|{func}|$".encode()
+                        await client.write_gatt_char(write_char, cmd)
+                        await asyncio.sleep(0.2)
+
+                        if responses:
+                            resp = responses[0].decode('ascii', errors='replace')
+                            self.log(f"    CMD|{func}|$ -> {resp}")
+                            self.report.working_commands.append(f"CMD|{func}|$")
+
+                    await client.stop_notify(notify_char)
+
+        # Summary
+        self.log("\n" + "=" * 60)
+        self.log("ENUMERATION COMPLETE")
+        self.log("=" * 60)
+        self.log(f"\nDevice: {self.report.device_name}")
+        self.log(f"Address: {self.report.device_address}")
+        self.log(f"Services: {len(self.report.services)}")
+        self.log(f"Working protocols: {[p.name for p in self.report.protocols if p.working]}")
+        self.log(f"Working commands: {len(self.report.working_commands)}")
+
+        return self.report
+
+
+async def enumerate_ble_device(
+    address: str,
+    output_path: Optional[Path] = None,
+    verbose: bool = True,
+) -> BLEEnumerationReport:
+    """
+    Convenience function to enumerate a BLE device.
+
+    Args:
+        address: BLE device address or name
+        output_path: Optional path to save JSON report
+        verbose: Print progress information
+
+    Returns:
+        BLEEnumerationReport with all discovered information
+    """
+    enumerator = BLEEnumerator(address, verbose=verbose)
+    report = await enumerator.enumerate()
+
+    if output_path:
+        report.to_json(output_path)
+        if verbose:
+            print(f"\nReport saved to: {output_path}")
+
+    return report
+
+
+# CLI integration
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python ble_enumerate.py <device_address> [output.json]")
+        print("\nExample:")
+        print("  python ble_enumerate.py FAF198B6-D9A4-CC27-7BBA-3159F63A61A9")
+        print("  python ble_enumerate.py mechdog_00 enumeration.json")
+        sys.exit(1)
+
+    address = sys.argv[1]
+    output = Path(sys.argv[2]) if len(sys.argv) > 2 else None
+
+    asyncio.run(enumerate_ble_device(address, output))