488 lines
21 KiB
Python
488 lines
21 KiB
Python
"""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
|