Initial public release
This commit is contained in:
+547
@@ -0,0 +1,547 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user