Initial public release

ble_keystore.py

@@ -0,0 +1,158 @@
+"""Persistent storage for the 32-byte BLE payload-encryption key(s).
+
+Each M5Stack generates its own key on first boot (per-device, persisted
+to LittleFS). The host pulls the device's key during every profile
+upload and stores it locally so subsequent BLE frames can be decrypted.
+
+Storage:
+    - ``config/.blekeyfile``        — legacy single-key file. The most
+      recently uploaded key. Kept as a fallback for situations where a
+      frame's device tag isn't in the per-MAC store.
+    - ``config/.ble_keys.json``     — per-MAC key store, keyed by the
+      eFuse base MAC (the same MAC that appears in each device's tag
+      string "M5Stack|AA:BB:CC:DD:EE:FF"). Multiple ATOMS3s can coexist
+      here so the user doesn't have to re-upload to switch devices.
+
+Public API (back-compat preserved):
+    load_key()                 -> bytes | None   (legacy file)
+    save_key(key)              -> None           (writes legacy file)
+    save_key_for_mac(mac, key) -> None           (writes per-MAC store)
+    load_key_for_mac(mac)      -> bytes | None
+    all_known_macs()           -> list[str]
+"""
+
+import json
+import os
+import threading
+from utils.constants import APPDATA_DIR
+
+KEY_LEN = 32
+_KEY_FILE = os.path.join(APPDATA_DIR, ".blekeyfile")
+_PER_MAC_FILE = os.path.join(APPDATA_DIR, ".ble_keys.json")
+
+_lock = threading.Lock()
+
+
+def key_path() -> str:
+    return _KEY_FILE
+
+
+# ---- Legacy single-key API ----
+
+def load_key() -> bytes | None:
+    try:
+        with open(_KEY_FILE, "rb") as f:
+            data = f.read()
+    except OSError:
+        return None
+    return data if len(data) == KEY_LEN else None
+
+
+def save_key(key: bytes) -> None:
+    if len(key) != KEY_LEN:
+        raise ValueError(f"BLE key must be {KEY_LEN} bytes, got {len(key)}")
+    os.makedirs(APPDATA_DIR, exist_ok=True)
+
+    # Detect a key change so we can wipe replay-protection counters: a new
+    # key invalidates any captured ciphertext, so old counters carry no
+    # protection value AND would block legitimate frames from a re-flashed
+    # device until we manually reset.
+    previous = load_key()
+
+    tmp = _KEY_FILE + ".tmp"
+    with open(tmp, "wb") as f:
+        f.write(key)
+    os.replace(tmp, _KEY_FILE)
+
+    if previous != key:
+        try:
+            import ble_replay
+            ble_replay.ReplayState().reset()
+        except Exception:
+            # Best-effort — replay reset failure isn't fatal (worst case,
+            # the next BLE frame from the device gets rejected and the user
+            # has to re-upload or wipe).
+            pass
+
+
+# ---- Per-MAC keystore ----
+
+def _normalize_mac(mac: str) -> str:
+    """Canonicalize MAC string: uppercase, colon-separated. Accepts the
+    device tag "M5Stack|AA:BB:..." OR a bare MAC."""
+    if mac.startswith("M5Stack|"):
+        mac = mac[len("M5Stack|"):]
+    return mac.upper()
+
+
+def _load_per_mac_locked() -> dict:
+    try:
+        with open(_PER_MAC_FILE, "r", encoding="utf-8") as f:
+            d = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return {}
+    if not isinstance(d, dict):
+        return {}
+    return d
+
+
+def _save_per_mac_locked(d: dict) -> None:
+    os.makedirs(APPDATA_DIR, exist_ok=True)
+    tmp = _PER_MAC_FILE + ".tmp"
+    with open(tmp, "w", encoding="utf-8") as f:
+        json.dump(d, f, indent=2)
+    os.replace(tmp, _PER_MAC_FILE)
+
+
+def save_key_for_mac(mac: str, key: bytes) -> None:
+    """Persist ``key`` (32 raw bytes) under ``mac``. Replaces any
+    existing key for that MAC and resets replay counters if the key
+    actually changed."""
+    if len(key) != KEY_LEN:
+        raise ValueError(f"BLE key must be {KEY_LEN} bytes, got {len(key)}")
+    mac = _normalize_mac(mac)
+    with _lock:
+        store = _load_per_mac_locked()
+        prev_hex = store.get(mac)
+        new_hex = key.hex()
+        if prev_hex == new_hex:
+            return
+        store[mac] = new_hex
+        _save_per_mac_locked(store)
+    try:
+        import ble_replay
+        ble_replay.ReplayState().reset()
+    except Exception:
+        pass
+
+
+def load_key_for_mac(mac: str) -> bytes | None:
+    """Look up the key for ``mac`` (accepts bare MAC or full tag).
+    Returns None if no key is stored for this device."""
+    mac = _normalize_mac(mac)
+    with _lock:
+        store = _load_per_mac_locked()
+    hex_key = store.get(mac)
+    if not hex_key:
+        return None
+    try:
+        b = bytes.fromhex(hex_key)
+    except ValueError:
+        return None
+    return b if len(b) == KEY_LEN else None
+
+
+def load_key_for_mac_or_default(mac: str) -> bytes | None:
+    """Per-MAC lookup with the legacy single-key file as a fallback.
+    Use this in the BLE layer when receiving a frame: try the right
+    key for the device's MAC first, fall back to the legacy file for
+    users who haven't re-uploaded since multi-device support landed."""
+    k = load_key_for_mac(mac)
+    if k is not None:
+        return k
+    return load_key()
+
+
+def all_known_macs() -> list:
+    with _lock:
+        return sorted(_load_per_mac_locked().keys())