"""ESP-NOW mesh keyboard manager — drop-in for MultiBleKeyboardManager. Replaces the per-device BLE fan-out with a single broadcast through the USB-attached hub device. Exposes the SAME public surface the BT Keyboard window already calls (slots / add_device / remove_device / set_enabled / set_label / identify / send_event / send_event_with_t / send_mouse / reset_session_clocks / stats / shutdown / set_callbacks), so the window's streaming, recording, and replay paths are unchanged. How it differs from the BLE manager internally: * One session per window. On open we mint a random 64-bit session_id and a fresh 32-byte AES group key. Every reliable frame the host emits carries a monotonically increasing transport seq; the hub caches it and retransmits on a node's NACK, so no keystroke is ever lost. * A "slot" is a known device MAC (the WiFi STA MAC == the AES device-tag MAC == the key in config/.ble_keys.json — no migration needed). We JOIN it by sending an encrypted invite (its per-device key) carrying the group key + index + base seq; the node decrypts, adopts the group key, and starts ACKing. * Discovery surfaces beacons the hub forwards (board type + label), independent of Bleak. * Status is derived from the hub's periodic ack-table: a node whose ack is fresh and whose cum is near the head is "connected"; lagging / gone nodes show accordingly. Threading: public API is Tk-thread safe. Outgoing keystrokes are batched on a background flush thread (≤16 events per frame, every few ms) so a burst pays one broadcast instead of N. Inbound hub callbacks fire on the MeshLink reader thread; we marshal state changes out via the same on_status_change / on_stats_change callbacks the window already installs. """ from __future__ import annotations import os import threading import time from dataclasses import dataclass, field from typing import Callable, Optional import ble_keystore from ble_frame import build_frame, DEVICE_TAG_PREFIX from live_protocol import ( MSG_STOP, MSG_KEYS, MSG_MOUSE, MSG_IDENTIFY, MSG_LABEL, MSG_PAUSE, MSG_RESUME, MESH_T_DATA, MESH_T_DATA_U, MESH_T_JOIN, MESH_T_BEACON, MESH_T_JOIN_ACK, MESH_T_ERR, BCAST_MAC, pack_header, pack_keys_body, pack_mouse_body, pack_mesh_header, parse_mesh_header, mac_to_bytes, ERR_LABELS, ) # Status strings (kept identical to ble_live's so the window's color map # and any existing checks keep working). ST_DISCONNECTED = "disconnected" ST_CONNECTING = "connecting" # JOIN sent, no ack yet ST_CONNECTED = "connected" ST_LAGGING = "lagging" ST_ERROR = "error" # Soft warning threshold for the ESP-NOW hub. Higher than BLE mode # (ble_multi.MAX_SLOTS == 4) because the hub broadcasts one frame to the # whole fleet instead of holding N concurrent BLE links. Not a hard cap — # the hub roster can track more; the UI just warns past this. MAX_SLOTS = 12 # A node is considered present if its last ack-table entry is younger than # this; lagging if older but still in the table. _ACK_FRESH_MS = 1500 _ACK_LAG_MS = 600 # cum-lag (frames behind head) that flips lagging _JOIN_RETRY_S = 1.0 # re-JOIN an unacked / offline node this often _FLUSH_INTERVAL_S = 0.005 # keystroke batch flush cadence _MAX_BATCH = 16 def _device_tag(mac: str) -> str: return DEVICE_TAG_PREFIX + mac.upper() @dataclass class DeviceSlot: address: str # STA MAC ("AA:BB:CC:DD:EE:FF") label: str = "" enabled: bool = True status: str = ST_DISCONNECTED board: str = "" # "atoms3" / "atoms3_lite" last_status_change: float = field(default_factory=time.monotonic) added_at: float = field(default_factory=time.monotonic) # Reliability bookkeeping (updated from the hub ack-table) cum: int = 0 ack_lag: int = 0 last_ack_age_ms: int = 0 events_sent: int = 0 bytes_sent: int = 0 joined: bool = False last_join_s: float = 0.0 idx: int = 0 def display_label(self) -> str: return self.label or self.address class MeshKeyboardManager: """Owns the mesh session and every known device slot.""" def __init__(self, link, max_slots: int = MAX_SLOTS): self._link = link self._max_slots = max_slots self._slots: dict[str, DeviceSlot] = {} self._lock = threading.Lock() self._on_status_change: Optional[Callable[[str, str], None]] = None self._on_stats_change: Optional[Callable[[], None]] = None # Session identity. Random session id + group key; no persistence — # a fresh window session can never collide with a stale one, and a # rebooted node re-JOINs into the current generation. self._session_id = int.from_bytes(os.urandom(8), "little") self._group_key = os.urandom(32) self._hub_mac = "" # filled from the espnow_hub response self._hub_tag = "" # Reliable transport seq (DATA lane) and move seq (DATA_U lane). self._seq = 0 self._move_seq = 0 self._seq_lock = threading.Lock() # Host session clock (mirrors ble_live: t=0 at first send). self._session_t0_ns: Optional[int] = None # Outgoing keystroke batch queue (flush thread coalesces). self._pending: list = [] self._pending_lock = threading.Lock() self._running = True self._flush_thread = threading.Thread( target=self._flush_loop, daemon=True, name="MeshFlush") # Discovery: beacons collected between discover() calls. self._beacons: dict[str, dict] = {} self._link.start(on_rx=self._on_hub_rx, on_acktab=self._on_acktab, on_json=self._on_hub_json) self._flush_thread.start() # ---- Hub identity ---- def set_hub(self, mac: str) -> None: """Record which device is the hub so the host never tries to JOIN it as a node, and so the data-lane tag binds to it.""" self._hub_mac = (mac or "").upper() self._hub_tag = _device_tag(self._hub_mac) if self._hub_mac else "" def hub_mac(self) -> str: return self._hub_mac # ---- Public API (parity with MultiBleKeyboardManager) ---- def set_callbacks(self, *, on_status_change=None, on_stats_change=None): self._on_status_change = on_status_change self._on_stats_change = on_stats_change def slots(self) -> list[DeviceSlot]: with self._lock: return sorted(self._slots.values(), key=lambda s: s.added_at) def slot_count(self) -> int: with self._lock: return len(self._slots) def is_full(self) -> bool: return False def over_soft_limit(self) -> bool: return self.slot_count() > self._max_slots def add_device(self, address: str, label: str = "", board: str = "") -> DeviceSlot | None: addr = address.upper() if self._hub_mac and addr == self._hub_mac: # The hub bridges; it is not a controllable node. return None with self._lock: if addr in self._slots: return None idx = len(self._slots) slot = DeviceSlot(address=addr, label=label or addr, board=board, idx=idx, status=ST_CONNECTING) self._slots[addr] = slot # Kick off a JOIN immediately; the flush/maintenance loop re-tries. self._send_join(slot) self._emit_status(addr, ST_CONNECTING) return slot def remove_device(self, address: str) -> bool: addr = address.upper() with self._lock: slot = self._slots.pop(addr, None) if slot is None: return False # Tell the node to leave so it stops emitting / releases held keys. self._send_inner(MSG_STOP, b"", dest=addr, reliable=True) return True def set_enabled(self, address: str, enabled: bool) -> None: addr = address.upper() with self._lock: slot = self._slots.get(addr) if slot is None: return slot.enabled = bool(enabled) # PAUSE/RESUME keeps the node ACKing (instant re-enable, no flood) # but stops it feeding its USB HID. self._send_inner(MSG_RESUME if enabled else MSG_PAUSE, b"", dest=addr, reliable=True) def get_slot(self, address: str) -> DeviceSlot | None: with self._lock: return self._slots.get(address.upper()) def set_label(self, address: str, label: str) -> None: addr = address.upper() with self._lock: slot = self._slots.get(addr) if slot is None: return slot.label = label or slot.address body = (label or "").encode("utf-8")[:38] self._send_inner(MSG_LABEL, body, dest=addr, reliable=True) def identify(self, address: str, on: bool = True) -> None: self._send_inner(MSG_IDENTIFY, bytes([1 if on else 0]), dest=address.upper(), reliable=True) # ---- Streaming ---- def send_event(self, action: int, hid_code: int) -> int: now_ns = time.monotonic_ns() if self._session_t0_ns is None: self._session_t0_ns = now_ns t_ms = (now_ns - self._session_t0_ns) // 1_000_000 if t_ms > 0xFFFFFFFF: t_ms = 0xFFFFFFFF return self.send_event_with_t(action, hid_code, int(t_ms)) def send_event_with_t(self, action: int, hid_code: int, t_ms: int) -> int: if not self._enabled_count(): return 0 with self._pending_lock: self._pending.append((int(action) & 0xFF, int(hid_code) & 0xFF, int(t_ms) & 0xFFFFFFFF)) return self._enabled_count() def send_mouse(self, buttons: int, x: float, y: float, wheel: int = 0) -> int: n = self._enabled_count() if not n: return 0 xi = int(max(0.0, min(1.0, x)) * 32767) yi = int(max(0.0, min(1.0, y)) * 32767) body = pack_mouse_body(int(buttons) & 0xFF, xi, yi, int(wheel)) # Button changes and wheel ticks must not be lost → reliable DATA; # pure moves ride the latest-wins DATA_U lane. reliable = (buttons != self._last_mouse_buttons) or (int(wheel) != 0) self._last_mouse_buttons = buttons self._send_inner(MSG_MOUSE, body, dest=BCAST_MAC, reliable=reliable) return n _last_mouse_buttons = 0 def reset_session_clocks(self) -> None: self._session_t0_ns = None # ---- Stats ---- def stats(self) -> list[dict]: out = [] for slot in self.slots(): out.append({ "address": slot.address, "label": slot.display_label(), "enabled": slot.enabled, "status": slot.status, "events_sent": slot.events_sent, "bytes_sent": slot.bytes_sent, "ack_lag": slot.ack_lag, "board": slot.board, }) return out def take_beacons(self) -> list[dict]: """Return discovered (non-slot) devices seen since the last call.""" with self._lock: known = set(self._slots.keys()) found = [v for k, v in self._beacons.items() if k not in known] self._beacons.clear() return found # ---- Shutdown ---- def shutdown(self) -> None: self._running = False # Best-effort: tell every node to leave the session. try: self._send_inner(MSG_STOP, b"", dest=BCAST_MAC, reliable=True) except Exception: pass with self._lock: self._slots.clear() if self._flush_thread.is_alive(): self._flush_thread.join(timeout=1.0) # ================================================================== # Internal — framing & send # ================================================================== def _next_seq(self) -> int: with self._seq_lock: self._seq += 1 return self._seq def _next_move_seq(self) -> int: with self._seq_lock: self._move_seq += 1 return self._move_seq def _build_data_frame(self, mtype: int, seq: int, dest: bytes, inner_plain: bytes) -> bytes | None: if not self._hub_tag: return None try: enc = build_frame(self._group_key, self._hub_tag, inner_plain) except Exception: return None return pack_mesh_header(mtype, 0, seq, dest) + enc def _send_inner(self, msg_type: int, body: bytes, dest, reliable: bool) -> None: """Encrypt one inner live-protocol message under the group key and ship it to the hub for broadcast. ``dest`` is BCAST_MAC for everyone or a MAC string for a single node.""" if isinstance(dest, str): dest_bytes = mac_to_bytes(dest) else: dest_bytes = dest if reliable: seq = self._next_seq() mtype = MESH_T_DATA else: seq = self._next_move_seq() mtype = MESH_T_DATA_U plain = pack_header(msg_type, self._session_id, seq) + body frame = self._build_data_frame(mtype, seq, dest_bytes, plain) if frame is None: return self._link.send_mesh_frame(frame) def _send_keys_batch(self, batch: list) -> None: if not batch: return seq = self._next_seq() body = pack_keys_body(batch) plain = pack_header(MSG_KEYS, self._session_id, seq) + body frame = self._build_data_frame(MESH_T_DATA, seq, BCAST_MAC, plain) if frame is None: return if self._link.send_mesh_frame(frame): with self._lock: for slot in self._slots.values(): if slot.enabled: slot.events_sent += len(batch) slot.bytes_sent += len(frame) def _send_join(self, slot: DeviceSlot) -> None: """Invite one device into the session: an AES frame under ITS per-device key carrying the group key, index, and current base seq (so a mid-session join doesn't trigger a historic NACK storm).""" key = ble_keystore.load_key_for_mac_or_default(slot.address) if key is None: self._emit_status(slot.address, ST_ERROR) return import json as _json base = self._seq payload = _json.dumps({ "sid": self._session_id, "gkey": self._group_key.hex(), "idx": slot.idx, "base": base, "label": slot.label, }).encode("utf-8") try: enc = build_frame(key, _device_tag(slot.address), payload) except Exception: self._emit_status(slot.address, ST_ERROR) return frame = pack_mesh_header(MESH_T_JOIN, 0, base, mac_to_bytes(slot.address)) + enc slot.last_join_s = time.monotonic() self._link.send_mesh_frame(frame) # Re-send the latched label so a freshly-joined node shows it. if slot.label and slot.label != slot.address: self._send_inner(MSG_LABEL, slot.label.encode("utf-8")[:38], dest=slot.address, reliable=True) def _enabled_count(self) -> int: with self._lock: return sum(1 for s in self._slots.values() if s.enabled) # ================================================================== # Internal — flush loop & maintenance # ================================================================== def _flush_loop(self) -> None: last_maint = 0.0 while self._running: # Drain pending keystrokes into ≤16-event batches. batch = None with self._pending_lock: if self._pending: batch = self._pending[:_MAX_BATCH] del self._pending[:_MAX_BATCH] if batch: self._send_keys_batch(batch) # If more remain, loop again immediately (no sleep). with self._pending_lock: if self._pending: continue now = time.monotonic() if now - last_maint >= _JOIN_RETRY_S: last_maint = now self._maintain_joins(now) time.sleep(_FLUSH_INTERVAL_S) def _maintain_joins(self, now: float) -> None: """Re-JOIN any slot that hasn't acked yet or has gone offline.""" for slot in self.slots(): if slot.joined and slot.status == ST_CONNECTED: continue if (now - slot.last_join_s) >= _JOIN_RETRY_S: self._send_join(slot) # ================================================================== # Internal — inbound from hub # ================================================================== def _on_acktab(self, entries: list) -> None: head = self._seq changed = False seen = set() for e in entries: mac = e["mac"].upper() seen.add(mac) with self._lock: slot = self._slots.get(mac) if slot is None: continue slot.cum = e["cum"] slot.ack_lag = max(0, head - e["cum"]) slot.last_ack_age_ms = e["age_ms"] fresh = e["age_ms"] < _ACK_FRESH_MS if fresh: slot.joined = True new_status = (ST_CONNECTED if slot.ack_lag <= _ACK_LAG_MS else ST_LAGGING) else: new_status = ST_LAGGING if new_status != slot.status: slot.status = new_status slot.last_status_change = time.monotonic() changed = True # Slots not in the table at all → disconnected. with self._lock: for mac, slot in self._slots.items(): if mac not in seen and slot.status not in (ST_CONNECTING, ST_DISCONNECTED): if slot.last_ack_age_ms or slot.joined: slot.status = ST_DISCONNECTED slot.joined = False changed = True if changed and self._on_status_change: try: self._on_status_change("", "acktab") except Exception: pass def _on_hub_rx(self, src_mac: str, frame: bytes) -> None: parsed = parse_mesh_header(frame) if parsed is None: return mtype, flags, seq, dest, payload = parsed if mtype == MESH_T_BEACON: self._handle_beacon(src_mac, payload) elif mtype == MESH_T_JOIN_ACK: self._handle_join_ack(src_mac, payload) elif mtype == MESH_T_ERR: self._handle_err(src_mac, payload) def _handle_beacon(self, src_mac: str, payload: bytes) -> None: # payload: ver, board(0/1), in_session, paused, label[...] board = "" label = "" if len(payload) >= 4: board = "atoms3_lite" if payload[1] == 1 else "atoms3" label = payload[4:].split(b"\x00", 1)[0].decode("utf-8", "ignore") with self._lock: self._beacons[src_mac.upper()] = { "address": src_mac.upper(), "board": board, "label": label, "name": label or None, } def _handle_join_ack(self, src_mac: str, payload: bytes) -> None: addr = src_mac.upper() with self._lock: slot = self._slots.get(addr) if slot is None: return slot.joined = True if slot.status == ST_CONNECTING: slot.status = ST_CONNECTED slot.last_status_change = time.monotonic() self._emit_status(addr, ST_CONNECTED) def _handle_err(self, src_mac: str, payload: bytes) -> None: if not payload: return code = payload[0] label = ERR_LABELS.get(code, f"ERR_{code}") if self._on_status_change: try: self._on_status_change(src_mac.upper(), f"error:{label}") except Exception: pass def _on_hub_json(self, doc: dict) -> None: # hub_pong / ok responses — currently informational only. pass def _emit_status(self, addr: str, status: str) -> None: if self._on_status_change: try: self._on_status_change(addr, status) except Exception: pass