"""BLE client for two-way encrypted variable sync with the ATOMS3. Architecture: - ESP32 AtomS3 = BLE GATT Server (peripheral), advertises on-demand when a Variables node in BLE mode (pull / push / request) is hit. - Python app = BLE Client (central, using bleak), continuously scans, connects when a device appears, subscribes to its notify characteristic. Wire framing (every BLE message, both directions): byte tag_len 1 byte bytes tag ASCII "M5Stack|AA:BB:CC:DD:EE:FF" bytes nonce 12 bytes bytes ciphertext+gcm_tag N + 16 bytes (AES-256-GCM, AAD = tag bytes) Plaintext is JSON with an `op` discriminator: {"op":"hello","kind":"pull"|"push"|"request","scope":"device"|"universal","names":[...]} Device's first notify after connect — declares what it wants. {"op":"pull","scope":"device"|"universal","vars":{...}} Host -> device. Variables to install in the chosen on-device store. {"op":"push","vars":{...}} Device -> host. {"op":"request","names":[...]} Device -> host. Followed by host pull. {"op":"ack"} Either direction. End of exchange. Frames whose tag prefix isn't "M5Stack|" or whose AEAD authentication fails are dropped silently — that's the "encrypted or thrown out" rule. """ import asyncio import json import logging import threading import ble_keystore import ble_replay import ble_debug_log as _bled from ble_frame import build_frame as _build_frame, parse_frame as _parse_frame from ble_frame import parse_frame_auto as _parse_frame_auto from ble_frame import DEVICE_TAG_PREFIX, MAC_RE # re-exported for callers # noqa: F401 log = logging.getLogger(__name__) SERVICE_UUID = "4fafc201-1fb5-459e-8fcc-c5c9c331914b" VARS_WRITE_UUID = "beb5483e-36e1-4688-b7f5-ea07361b26a8" # host -> device (write) VARS_NOTIFY_UUID = "beb5483e-36e1-4688-b7f5-ea07361b26a9" # device -> host (notify) # Live keystroke streaming channel — see ble_live.py for the protocol. # WRITE characteristic is Write-Without-Response for low latency. LIVE_KEYS_WRITE_UUID = "4fafc202-1fb5-459e-8fcc-c5c9c331914b" # host -> device (WWR) LIVE_KEYS_NOTIFY_UUID = "4fafc203-1fb5-459e-8fcc-c5c9c331914b" # device -> host (notify) # Advertised service UUID specifically for live mode. Disjoint from # SERVICE_UUID so the var-sync scan filter and the live-keystroke scan # filter never resolve to the same device — eliminates the race where # both clients try to grab the BLE connection at once. LIVE_SERVICE_UUID = "4fafc204-1fb5-459e-8fcc-c5c9c331914b" SCAN_TIMEOUT_S = 3 SCAN_PAUSE_S = 1 # How long to wait for the device's first frame (hello) after subscription. # Bleak + NimBLE on ESP32-S3 with concurrent USB-CDC sometimes takes 10+ # seconds to actually deliver a notification, so this is intentionally # loose. The device-side timeout is what really gates the exchange. EXCHANGE_TIMEOUT_S = 30 class BLEVariableClient: """BLE GATT client — scans, connects, exchanges encrypted op frames.""" def __init__(self): self._running = False self._thread: threading.Thread | None = None self._loop: asyncio.AbstractEventLoop | None = None self._on_status = None # Caller-supplied callbacks (all invoked on the BLE thread): # get_vars(scope, mac) -> dict # set_device_vars(mac, vars_dict) — store push'd dict # prompt_request(mac, names) -> dict|None — show modal in GUI thread, # return user's edits or None # on_device_seen(mac) — first time we see this MAC self._get_vars = None self._set_device_vars = None self._prompt_request = None self._on_device_seen = None self._seen_macs: set[str] = set() # Cached project ble_variables dict — used as a fallback when # `_get_vars` callback isn't wired yet (early startup). self._cached_vars = {"universal": {}, "devices": {}} # Replay-protection counter store. Persists per-host send counter # and per-device last-seen counter to disk. self._replay = ble_replay.ReplayState() # Public API (Tk thread) def start(self, ble_variables, on_status=None, get_vars=None, set_device_vars=None, prompt_request=None, on_device_seen=None) -> None: self._cached_vars = self._normalize(ble_variables) self._on_status = on_status self._get_vars = get_vars self._set_device_vars = set_device_vars self._prompt_request = prompt_request self._on_device_seen = on_device_seen self._running = True _bled.event("session_start", cached_universal=len(self._cached_vars.get("universal") or {}), cached_devices=len(self._cached_vars.get("devices") or {})) self._thread = threading.Thread(target=self._thread_main, daemon=True, name="BLEVariableClient") self._thread.start() def update_variables(self, ble_variables) -> None: self._cached_vars = self._normalize(ble_variables) def stop(self) -> None: self._running = False if self._loop and self._loop.is_running(): self._loop.call_soon_threadsafe(self._loop.stop) if self._thread and self._thread.is_alive(): self._thread.join(timeout=4) @staticmethod def _normalize(raw) -> dict: if not isinstance(raw, dict): return {"universal": {}, "devices": {}} if "universal" in raw or "devices" in raw: return { "universal": dict(raw.get("universal") or {}), "devices": {str(k): dict(v or {}) for k, v in (raw.get("devices") or {}).items()}, } return {"universal": dict(raw), "devices": {}} def _vars_for(self, scope: str, mac: str) -> dict: if self._get_vars is not None: try: return dict(self._get_vars(scope, mac) or {}) except Exception as e: log.warning("get_vars callback error: %s", e) if scope == "device": return dict(self._cached_vars.get("devices", {}).get(mac, {})) return dict(self._cached_vars.get("universal", {})) # Internal — runs on the BLE thread / event loop def _thread_main(self) -> None: try: from bleak import BleakClient, BleakScanner # noqa except ImportError: print("[BLE] bleak not installed. Run: pip install bleak") self._set_status("error") return self._BleakClient = BleakClient self._BleakScanner = BleakScanner self._loop = asyncio.new_event_loop() asyncio.set_event_loop(self._loop) try: self._loop.run_until_complete(self._run()) except RuntimeError: pass except Exception as exc: log.warning("BLE client error: %s", exc, exc_info=True) finally: try: self._loop.close() except Exception: pass async def _run(self) -> None: while self._running: await self._scan_once() if self._running: await asyncio.sleep(SCAN_PAUSE_S) async def _scan_once(self) -> None: BleakScanner = self._BleakScanner BleakClient = self._BleakClient try: _bled.event("scan_start", timeout_s=SCAN_TIMEOUT_S) device = await BleakScanner.find_device_by_filter( lambda d, adv: SERVICE_UUID in (adv.service_uuids or []), timeout=SCAN_TIMEOUT_S, ) if not device: _bled.event("scan_idle") self._set_status("idle") return # Per-MAC keystore. We don't pre-load a single key here — # we don't know which device we connected to until we read # the tag from its hello frame. ``_exchange`` resolves the # per-MAC key on receipt of the first frame and uses it for # all subsequent encrypts/decrypts in this session. self._set_status("connecting") _bled.event("scan_found", name=device.name, address=device.address) print(f"[BLE Client] Connected to {device.name} ({device.address})") # Cap the connect itself. Without this, a device caught mid- # reboot (e.g. firmware panicked on the previous exchange) or a # wedged WinRT stack leaves us in `async with` indefinitely and # the toolbar stays stuck on "Connecting...". async with BleakClient(device, timeout=15.0) as client: _bled.event("connected", address=device.address) await self._exchange(client) _bled.event("exchange_returned", address=device.address) _bled.event("disconnected", address=device.address) self._set_status("synced") except asyncio.TimeoutError as exc: _bled.event("connect_timeout", error=repr(exc)) print(f"[BLE Client] Connect timed out: {exc}") self._set_status("error") except Exception as exc: _bled.event("scan_exception", error=repr(exc)) print(f"[BLE Client] Error: {exc}") self._set_status("error") async def _exchange(self, client) -> None: """One full exchange with a connected device. Subscribes to notifications, waits for the device's first frame (which declares its `kind` via op=hello), responds appropriately, then waits for either the device to disconnect or for a follow-up frame, then returns. Uses per-MAC key resolution: parse_frame_auto picks the right key from the keystore based on the tag in the frame, so a host with multiple devices' keys stored picks the right one automatically. ``session_key`` captures the resolved key for outgoing frames in this exchange. """ first_frame: asyncio.Future = self._loop.create_future() follow_frame: asyncio.Queue = asyncio.Queue() # Resolved on first successful decrypt; used for sends below. session_key_ref = {"key": None} def handle_notify(_char, data: bytearray): raw = bytes(data) result = _parse_frame_auto(raw) if result is None: _bled.event("notify_drop_malformed", bytes_len=len(raw)) return tag, plain = result if plain is None: _bled.event("notify_drop_auth", bytes_len=len(raw), tag=tag) return # Cache the resolved key for outgoing pull/ack frames. if session_key_ref["key"] is None: mac = tag[len(DEVICE_TAG_PREFIX):] session_key_ref["key"] = ble_keystore.load_key_for_mac_or_default(mac) parsed = (tag, plain) _bled.event( "notify_rx", bytes_len=len(raw), plain_len=len(plain), tag=tag, first=not first_frame.done(), ) if not first_frame.done(): first_frame.set_result(parsed) else: try: follow_frame.put_nowait(parsed) except Exception: pass try: await client.start_notify(VARS_NOTIFY_UUID, handle_notify) _bled.event("subscribed") except Exception as exc: _bled.event("start_notify_failed", error=repr(exc)) print(f"[BLE Client] start_notify failed: {exc}") return try: tag_str, plaintext = await asyncio.wait_for(first_frame, timeout=EXCHANGE_TIMEOUT_S) except asyncio.TimeoutError: _bled.event("hello_timeout", timeout_s=EXCHANGE_TIMEOUT_S) print("[BLE Client] Timed out waiting for device hello") return mac = tag_str[len(DEVICE_TAG_PREFIX):] # The per-MAC key for this device — resolved via the keystore # on the first successful decrypt (handle_notify above). Use # this for all outbound writes in this exchange so we encrypt # with the right key when the user has multiple devices on # file. key = session_key_ref.get("key") if key is None: _bled.event("no_key_for_mac", mac=mac) print(f"[BLE Client] No key on file for {mac} — upload the profile via USB first") return if mac not in self._seen_macs: self._seen_macs.add(mac) if self._on_device_seen: try: self._on_device_seen(mac) except Exception as e: log.warning("on_device_seen error: %s", e) try: msg = json.loads(plaintext.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: _bled.event("hello_bad_json", error=repr(exc)) return # Replay protection: every device frame carries (session_id, seq). # A new session_id signals a device reboot/reflash and resets the # window; same session_id requires strictly-increasing seq. seq = msg.get("seq") sid = msg.get("session_id") ok, reason = self._replay.accept_received(mac, sid, seq) if not ok: _bled.event("hello_replay_reject", mac=mac, seq=seq, session_id=sid, reason=reason) print(f"[BLE Client] dropped {reason} from {mac}: seq={seq} sid={sid}") return op = msg.get("op") _bled.event("hello_ok", mac=mac, op=op, hello_kind=msg.get("kind"), seq=seq, session_id=sid, accept_reason=reason) if op == "hello": kind = msg.get("kind", "pull") if kind == "pull": scope = msg.get("scope", "universal") if scope not in ("device", "universal"): scope = "universal" await self._send_pull(client, key, mac, scope, self._vars_for(scope, mac)) elif kind == "push": # Device immediately follows hello with the actual push frame. await self._handle_push_followup(client, key, mac, follow_frame) elif kind == "request": names = list(msg.get("names") or []) await self._handle_request(client, key, mac, names) else: _bled.event("hello_unknown_kind", kind=kind) print(f"[BLE Client] Unknown hello kind: {kind}") return # Some devices may skip the hello and send op=push directly. if op == "push": await self._apply_push(mac, msg) await self._send_ack(client, key, mac) return if op == "request": names = list(msg.get("names") or []) await self._handle_request(client, key, mac, names) return _bled.event("hello_unknown_op", op=op) print(f"[BLE Client] Unknown op: {op}") async def _handle_push_followup(self, client, key: bytes, mac: str, follow_frame: asyncio.Queue) -> None: try: tag_str, plaintext = await asyncio.wait_for( follow_frame.get(), timeout=EXCHANGE_TIMEOUT_S) except asyncio.TimeoutError: _bled.event("push_followup_timeout", mac=mac) print("[BLE Client] Timed out waiting for push payload") return if tag_str[len(DEVICE_TAG_PREFIX):] != mac: _bled.event("push_followup_mac_mismatch", expected=mac, got=tag_str) return try: msg = json.loads(plaintext.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): _bled.event("push_followup_bad_json") return seq = msg.get("seq") sid = msg.get("session_id") ok, reason = self._replay.accept_received(mac, sid, seq) if not ok: _bled.event("push_followup_replay_reject", mac=mac, seq=seq, session_id=sid, reason=reason) return if msg.get("op") != "push": _bled.event("push_followup_wrong_op", op=msg.get("op")) return _bled.event("push_received", mac=mac, seq=seq, session_id=sid, accept_reason=reason, var_count=len(msg.get("vars") or {})) await self._apply_push(mac, msg) await self._send_ack(client, key, mac) async def _apply_push(self, mac: str, msg: dict) -> None: vars_ = msg.get("vars") or {} if not isinstance(vars_, dict): return if self._set_device_vars: try: self._set_device_vars(mac, vars_) except Exception as e: log.warning("set_device_vars error: %s", e) # Mirror into cache so subsequent _vars_for("device", mac) matches. self._cached_vars.setdefault("devices", {})[mac] = dict(vars_) print(f"[BLE Client] Pushed {len(vars_)} vars from {mac}") async def _handle_request(self, client, key: bytes, mac: str, names: list) -> None: edited = None if self._prompt_request: # Tell the UI we're waiting on the user so the toolbar text is # distinguishable from a stuck-connect state. self._set_status("awaiting") # Run the blocking dialog wait in a thread so the asyncio loop # keeps processing notifications, disconnects, and the eventual # write_gatt_char back to the device. Calling the sync # prompt_request directly would block the entire BLE event loop # for as long as the dialog stays open. try: loop = asyncio.get_event_loop() edited = await asyncio.wait_for( loop.run_in_executor(None, self._prompt_request, mac, names), timeout=300, ) except asyncio.TimeoutError: _bled.event("prompt_request_timeout", mac=mac) log.warning("prompt_request timed out after 300s") edited = None except Exception as e: log.warning("prompt_request error: %s", e) edited = None if edited is None: # User cancelled or no callback wired — fall back to current values. current = self._vars_for("device", mac) edited = {n: current.get(n, "") for n in names} # Persist what the user entered so future Type Text expansions see it. if self._set_device_vars: current = self._vars_for("device", mac) current.update(edited) try: self._set_device_vars(mac, current) except Exception as e: log.warning("set_device_vars error: %s", e) self._cached_vars.setdefault("devices", {})[mac] = dict(current) await self._send_pull(client, key, mac, "device", edited) async def _send_pull(self, client, key: bytes, mac: str, scope: str, vars_: dict) -> None: tag = DEVICE_TAG_PREFIX + mac seq = self._replay.next_send_seq() sid = self._replay.host_session_id() plaintext = json.dumps({ "op": "pull", "scope": scope, "vars": vars_, "seq": seq, "session_id": sid, }).encode("utf-8") frame = _build_frame(key, tag, plaintext) _bled.event("send_pull", mac=mac, scope=scope, seq=seq, session_id=sid, var_count=len(vars_), frame_len=len(frame)) try: await client.write_gatt_char(VARS_WRITE_UUID, frame, response=True) _bled.event("send_pull_ok", mac=mac, seq=seq) print(f"[BLE Client] Sent pull(scope={scope}) to {mac}: {len(vars_)} vars") except Exception as exc: _bled.event("send_pull_failed", mac=mac, seq=seq, error=repr(exc)) print(f"[BLE Client] Pull write failed: {exc}") async def _send_ack(self, client, key: bytes, mac: str) -> None: tag = DEVICE_TAG_PREFIX + mac seq = self._replay.next_send_seq() sid = self._replay.host_session_id() plaintext = json.dumps({ "op": "ack", "seq": seq, "session_id": sid, }).encode("utf-8") frame = _build_frame(key, tag, plaintext) _bled.event("send_ack", mac=mac, seq=seq, session_id=sid, frame_len=len(frame)) try: await client.write_gatt_char(VARS_WRITE_UUID, frame, response=True) _bled.event("send_ack_ok", mac=mac, seq=seq) except Exception as exc: _bled.event("send_ack_failed", mac=mac, seq=seq, error=repr(exc)) print(f"[BLE Client] Ack write failed: {exc}") def _set_status(self, status: str) -> None: if self._on_status: try: self._on_status(status) except Exception: pass