133 lines
5.3 KiB
Python
133 lines
5.3 KiB
Python
"""Replay-protection state with per-session resilience.
|
|
|
|
Each side embeds (session_id, seq) in every frame. The session_id is a
|
|
random 64-bit value generated at boot/startup and stays constant for the
|
|
life of the process. The receiver tracks (last_session_id, max_seq) per
|
|
peer; a fresh session_id resets seq tracking gracefully (handles device
|
|
reflash, host wipe, etc.), while an unchanged session_id requires
|
|
strictly increasing seq (replay protection).
|
|
|
|
Tradeoff: an attacker who captures frames from an old session can replay
|
|
them AFTER the receiver has accepted a different session_id from the same
|
|
peer. For the user's threat model (passive sniffer in a controlled room,
|
|
no physical device access) this is acceptable. For stronger protection,
|
|
extend `device_state` to store a set of all ever-seen session_ids.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import threading
|
|
import secrets
|
|
from pathlib import Path
|
|
|
|
from utils.constants import APPDATA_DIR
|
|
|
|
_REPLAY_STATE_PATH = Path(APPDATA_DIR) / ".ble_replay.json"
|
|
|
|
|
|
def _new_session_id() -> int:
|
|
return secrets.randbits(64)
|
|
|
|
|
|
class ReplayState:
|
|
"""Per-MAC (session_id, max_seq) store + host's own session+counter."""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
# Host's outgoing identity. `host_session_id` is regenerated on
|
|
# every host restart so the device knows the host has reset and
|
|
# accepts the (new session, fresh seq) gracefully.
|
|
self._host_session_id: int = _new_session_id()
|
|
self._host_send_seq: int = 0
|
|
# Per-device tracking: {mac: {"session_id": int, "max_seq": int}}
|
|
self._device_state: dict[str, dict] = {}
|
|
self._load()
|
|
|
|
def _load(self) -> None:
|
|
try:
|
|
with open(_REPLAY_STATE_PATH, "r", encoding="utf-8") as f:
|
|
d = json.load(f)
|
|
# We DELIBERATELY do NOT restore host_session_id from disk:
|
|
# treating each Python process startup as a new session
|
|
# automatically heals "host has stale state" cases.
|
|
# We do persist host_send_seq within a session so a quick
|
|
# crash-restart in the same process doesn't reuse seqs (but
|
|
# since session_id changed, reuse is harmless anyway).
|
|
raw_dev = d.get("device_state") or {}
|
|
for mac, st in raw_dev.items():
|
|
if isinstance(st, dict):
|
|
self._device_state[str(mac)] = {
|
|
"session_id": int(st.get("session_id", 0)),
|
|
"max_seq": int(st.get("max_seq", 0)),
|
|
}
|
|
# Legacy migration: pre-session_id schema had `device_seen`.
|
|
# We can't recover the old session_id (it didn't exist), so
|
|
# we drop it — first frame from each device will be accepted
|
|
# fresh under the new session_id.
|
|
except (FileNotFoundError, json.JSONDecodeError, ValueError, TypeError):
|
|
pass
|
|
|
|
def _save_locked(self) -> None:
|
|
try:
|
|
_REPLAY_STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = _REPLAY_STATE_PATH.with_suffix(".tmp")
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump({
|
|
"host_send_seq": self._host_send_seq,
|
|
"device_state": self._device_state,
|
|
}, f)
|
|
os.replace(tmp, _REPLAY_STATE_PATH)
|
|
except Exception:
|
|
pass
|
|
|
|
def host_session_id(self) -> int:
|
|
return self._host_session_id
|
|
|
|
def next_send_seq(self) -> int:
|
|
"""Reserve the next outgoing seq within the current host session."""
|
|
with self._lock:
|
|
self._host_send_seq += 1
|
|
seq = self._host_send_seq
|
|
self._save_locked()
|
|
return seq
|
|
|
|
def accept_received(self, mac: str, session_id: int, seq: int) -> tuple[bool, str]:
|
|
"""Validate an inbound (session_id, seq) from ``mac``.
|
|
|
|
Returns (accepted, reason). ``reason`` is "fresh_session",
|
|
"monotonic", "regressed_seq", or "stale_session".
|
|
"""
|
|
if not isinstance(session_id, int) or not isinstance(seq, int):
|
|
return False, "bad_types"
|
|
with self._lock:
|
|
st = self._device_state.get(mac)
|
|
if st is None:
|
|
# Never-seen device — accept and remember.
|
|
self._device_state[mac] = {
|
|
"session_id": session_id,
|
|
"max_seq": seq,
|
|
}
|
|
self._save_locked()
|
|
return True, "fresh_device"
|
|
if st["session_id"] != session_id:
|
|
# Different session: device rebooted (or was reflashed).
|
|
# Accept fresh — this is the desync-recovery path.
|
|
st["session_id"] = session_id
|
|
st["max_seq"] = seq
|
|
self._save_locked()
|
|
return True, "fresh_session"
|
|
# Same session: seq must be strictly monotonic.
|
|
if seq <= st["max_seq"]:
|
|
return False, "regressed_seq"
|
|
st["max_seq"] = seq
|
|
self._save_locked()
|
|
return True, "monotonic"
|
|
|
|
def reset(self) -> None:
|
|
"""Wipe all counters. Call when the AES key changes."""
|
|
with self._lock:
|
|
self._host_session_id = _new_session_id()
|
|
self._host_send_seq = 0
|
|
self._device_state = {}
|
|
self._save_locked()
|