m5stack-automation-tool/CompileAndUpload/sync_key.py

"""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))