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