159 lines
4.9 KiB
Python
159 lines
4.9 KiB
Python
"""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())
|