"""Initial BLE key sync, run at the end of a firmware flash. A flash does a full chip erase (see upload.bat), which wipes the device's LittleFS — so the ATOMS3 generates a BRAND-NEW per-device AES-256 key on its first boot. Whatever key the host had stored is now stale, and live BLE recording / variable sync would fail GCM auth until the next profile upload re-synced it. This script closes that gap automatically right after flashing: it waits for the device to boot, pulls the fresh key over USB, and stores it in the host keystore (both the per-MAC entry and the legacy single-key file — exactly what serial_manager.upload_all does). Robustness: * The COM port commonly changes across the post-flash reboot, so we don't trust any hint — we scan every Espressif port each poll. * We wait up to BOOT_WAIT_S for the device to come up. If it hasn't appeared within REBOOT_HINT_S we print a one-time nudge to unplug/ replug (covers boards that don't auto-reboot out of download mode), then keep polling. * get_ble_key is retried a few times per session — the firmware needs LittleFS mounted and the keystore initialised, which finishes a beat into setup(). * Ports are opened with DTR/RTS low (the same non-resetting probe auto_enter_bootloader.py uses) and the firmware disables DTR-reboot, so connecting never knocks the device back into the bootloader. * Serial output is line-noisy (the firmware mirrors [BLE.dbg]/[BOOT] prints onto the same CDC pipe), so we skip non-JSON lines instead of failing on the first one. * Exit 0 on success, 1 otherwise — a failure is reported clearly but never marks the flash itself as failed (the .bat ignores our code). """ import json import os import sys import time # Make the repo root importable so `ble_keystore` (and its # `utils.constants` dependency) resolve when run from CompileAndUpload/. _REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) try: import serial import serial.tools.list_ports except ImportError: print(" ERROR: pyserial not installed — cannot sync the BLE key.", file=sys.stderr) sys.exit(1) import ble_keystore DEVICE_ID = "ATOMS3-MACROPAD" ESPRESSIF_VID = 0x303A # Total time to wait for the device to boot + answer. BOOT_WAIT_S = 40.0 # Initial settle before the first probe (USB re-enumeration after reset). INITIAL_SETTLE_S = 2.0 # If we still haven't seen the device by this point, nudge the user. REBOOT_HINT_S = 9.0 # Per-command read window. Generous because boot prints can precede the # JSON response on the shared CDC pipe. CMD_TIMEOUT_S = 3.0 # How many times to ask for the key once the device is confirmed. KEY_RETRIES = 5 def _espressif_ports() -> list: return [p for p in serial.tools.list_ports.comports() if p.vid == ESPRESSIF_VID] def _open(port: str): """Open the port without triggering a reset (DTR/RTS low, then raise DTR once the line is up — matches auto_enter_bootloader.py).""" ser = serial.Serial() ser.port = port ser.baudrate = 115200 ser.timeout = CMD_TIMEOUT_S ser.dtr = False ser.rts = False ser.open() time.sleep(0.1) ser.dtr = True time.sleep(0.3) ser.reset_input_buffer() return ser def _cmd(ser, obj: dict, expect_rsp: str) -> dict | None: """Send one JSON command and return the first reply whose "rsp" matches expect_rsp. Skips boot/debug noise lines.""" try: ser.reset_input_buffer() ser.write((json.dumps(obj) + "\n").encode("ascii")) ser.flush() except (serial.SerialException, OSError): return None deadline = time.monotonic() + CMD_TIMEOUT_S while time.monotonic() < deadline: try: line = ser.readline().decode("utf-8", errors="ignore").strip() except (serial.SerialException, OSError): return None if not line: continue try: d = json.loads(line) except json.JSONDecodeError: continue # firmware [BLE.dbg]/[BOOT] line — ignore if isinstance(d, dict) and d.get("rsp") == expect_rsp: return d return None def _try_port(port: str): """Confirm a MacroPad on `port` and pull its key. Returns (key_bytes, tag_or_None) on success, else None.""" try: ser = _open(port) except (serial.SerialException, OSError): return None try: ping = _cmd(ser, {"cmd": "ping"}, "pong") if not ping or ping.get("id") != DEVICE_ID: return None print(f" Found {DEVICE_ID} v{ping.get('ver', '?')} on {port}.") for attempt in range(KEY_RETRIES): rsp = _cmd(ser, {"cmd": "get_ble_key"}, "ble_key") if rsp: try: key = bytes.fromhex(rsp["key"]) except (KeyError, ValueError): key = None if key and len(key) == ble_keystore.KEY_LEN: tag = rsp.get("tag") return key, (tag if isinstance(tag, str) else None) time.sleep(1.0) print(" Device responded but did not return a valid BLE key.") return None finally: try: ser.close() except Exception: pass def _persist(key: bytes, tag: str | None) -> None: # Always update the legacy single-key file (back-compat), and the # per-MAC store when the device reported its tag. Both calls reset # replay-protection counters on a key change, which is exactly what we # want for a freshly-reflashed device (new key + new boot session). ble_keystore.save_key(key) if tag: ble_keystore.save_key_for_mac(tag, key) def main(argv: list) -> int: print("Syncing BLE key — waiting for the ATOMS3 to boot after flashing...") time.sleep(INITIAL_SETTLE_S) deadline = time.monotonic() + BOOT_WAIT_S nudged = False while time.monotonic() < deadline: for port in _espressif_ports(): res = _try_port(port.device) if res: key, tag = res try: _persist(key, tag) except OSError as exc: print(f" ERROR: could not write keystore: {exc}", file=sys.stderr) return 1 where = tag if tag else "(legacy key file — no device tag)" print(f" BLE key synced for {where}.") print(" Live BLE recording and variable sync are ready.") return 0 if not nudged and time.monotonic() > (deadline - BOOT_WAIT_S + REBOOT_HINT_S): print(" Still waiting... if the device did not reboot into the " "app on its own,") print(" unplug and replug it now (it will sync as soon as it " "boots).") nudged = True time.sleep(1.0) print(" Could not sync the BLE key (device never answered).") print(" No harm done — uploading a profile from the app syncs the key " "too. You can also re-run this step.") return 1 if __name__ == "__main__": sys.exit(main(sys.argv))