Save diagnostics to disk on connection error (#325)

Adds an opt-in option to automatically persist device diagnostics to disk when a BLE connection enters an error state, making it easier to troubleshoot intermittent connection / auth failures after the fact.

- New `with_diagnostics_on_exception(enabled)` builder on `DeviceBase` and matching `diagnostics_on_exception` toggle in the diagnostics options of the config flow. When enabled, packet collection is force-enabled and a connection-state listener saves the collected diagnostics to disk on any error state.
- Captures sent raw data and the initial session key in the diagnostics payload so post-mortem analysis has the full session context, not just received frames.
- Refactors `diagnostics.py` to build the dict via `DeviceDiagnosticsCollector.build_diagnostics_dict()`, consolidating the per-device fields in one place.

Author: GnoX

custom_components/ef_ble/__init__.py

@@ -29,6 +29,7 @@
     CONF_BLUEZ_START_NOTIFY,
     CONF_COLLECT_PACKETS_AMOUNT,
     CONF_CONNECTION_TIMEOUT,
+    CONF_DIAGNOSTICS_ON_EXCEPTION,
     CONF_DIAGNOSTICS_OPTIONS,
     CONF_EXTRA_BATTERY,
     CONF_PACKET_VERSION,
@@ -105,6 +106,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: DeviceConfigEntry) -> bo
     packet_collection_enabled = diag_options.get(
         CONF_COLLECT_PACKETS, eflib.is_unsupported(device)
     )
+    diagnostics_on_exception = diag_options.get(CONF_DIAGNOSTICS_ON_EXCEPTION, False)
 
     advanced = merged_options.get(CONF_ADVANCED_CONNECTION_OPTIONS, {})
     timeout = advanced.get(CONF_CONNECTION_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT)
@@ -121,6 +123,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: DeviceConfigEntry) -> bo
             .with_disabled_reconnect()
             .with_packet_version(packet_version.to_num())
             .with_enabled_packet_diagnostics(packet_collection_enabled)
+            .with_diagnostics_on_exception(diagnostics_on_exception)
             .with_connection_options(options)
             .connect(
                 user_id=user_id,
@@ -285,6 +288,7 @@ async def _update_listener(hass: HomeAssistant, entry: DeviceConfigEntry):
         CONF_COLLECT_PACKETS, eflib.is_unsupported(device)
     )
     diagnostics_buffer_size = diag_options.get(CONF_COLLECT_PACKETS_AMOUNT, 100)
+    diagnostics_on_exception = diag_options.get(CONF_DIAGNOSTICS_ON_EXCEPTION, False)
     advanced = merged_options.get(CONF_ADVANCED_CONNECTION_OPTIONS, {})
     options = Connection.Options(
         timeout=advanced.get(CONF_CONNECTION_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT),
@@ -298,5 +302,6 @@ async def _update_listener(hass: HomeAssistant, entry: DeviceConfigEntry):
             enabled=packet_collection,
             buffer_size=diagnostics_buffer_size,
         )
+        .with_diagnostics_on_exception(diagnostics_on_exception)
         .with_connection_options(options)
     )

custom_components/ef_ble/config_flow.py

@@ -42,6 +42,7 @@
     CONF_COLLECT_PACKETS_AMOUNT,
     CONF_CONNECTION_TIMEOUT,
     CONF_DIAGNOSTICS_ENCRYPT,
+    CONF_DIAGNOSTICS_ON_EXCEPTION,
     CONF_DIAGNOSTICS_OPTIONS,
     CONF_EXTRA_BATTERY,
     CONF_LOG_BLEAK,
@@ -433,10 +434,11 @@ async def _validate_user_id(
         if error := self._check_user_id(user_id):
             return error
 
-        device.with_logging_options(
-            ConfLogOptions.from_config(user_input)
-        ).with_packet_version(packet_version.to_num()).with_connection_options(
-            Connection.Options(timeout=timeout)
+        (
+            device.with_logging_options(ConfLogOptions.from_config(user_input))
+            .with_packet_version(packet_version.to_num())
+            .with_diagnostics_on_exception(True)
+            .with_connection_options(Connection.Options(timeout=timeout))
         )
 
         await device.connect(self._user_id)
@@ -445,7 +447,9 @@ async def _validate_user_id(
             conn_state, exc = await asyncio.wait_for(
                 device.wait_until_authenticated_or_error(return_exc=True), timeout
             )
-        except TimeoutError:
+        except TimeoutError as e:
+            exc = e
+            device.set_connection_state(ConnectionState.ERROR_TIMEOUT, e)
             conn_state = device.connection_state
 
         await device.disconnect()
@@ -707,6 +711,11 @@ def diagnostics_options(
                 vol.Required(CONF_DIAGNOSTICS_OPTIONS): section(
                     (
                         schema_builder()
+                        .optional(
+                            CONF_DIAGNOSTICS_ON_EXCEPTION,
+                            bool,
+                            diag.get(CONF_DIAGNOSTICS_ON_EXCEPTION, False),
+                        )
                         .optional(
                             CONF_COLLECT_PACKETS,
                             bool,

custom_components/ef_ble/const.py

@@ -16,6 +16,7 @@
 
 CONF_DIAGNOSTICS_OPTIONS = "diagnostics_options"
 CONF_DIAGNOSTICS_ENCRYPT = "diagnostics_encrypt"
+CONF_DIAGNOSTICS_ON_EXCEPTION = "diagnostics_on_exception"
 
 CONF_LOG_MASKED = "log_masked"
 CONF_LOG_PACKETS = "log_packets"

custom_components/ef_ble/diagnostics.py

@@ -14,29 +14,13 @@ async def async_get_config_entry_diagnostics(
 
     session = Session() if encrypt else None
 
-    diagnostics = {
-        "local_name": entry.data.get("local_name", None),
-        "device": device.device,
-        "name": device.name,
-        "default_name": device._default_name,
-        "sn_prefix": device._sn[:4],
-        "connection_state": device.connection_state,
-        "connection_state_history": list(device.connection_log.history),
-        "manufacturer_data": (
-            session.encrypt(device._manufacturer_data).hex()
-            if session is not None
-            else device._manufacturer_data.hex()
-        ),
-    }
-
-    if session is not None:
-        diagnostics["session"] = session.header.hex()
+    diagnostics: dict = {"local_name": entry.data.get("local_name", None)}
+    diagnostics |= device.diagnostics.build_diagnostics_dict(session)
 
     if device.diagnostics.is_enabled:
         connection_setup = await hass.async_add_executor_job(
             device.connection_log.load_from_cache
         )
-        diagnostics |= device.diagnostics.as_dict(session)
         diagnostics |= {"connection_setup": connection_setup}
 
     return diagnostics

custom_components/ef_ble/eflib/connection.py

@@ -229,6 +229,7 @@ def __init__(
         self._packet_version = packet_version
         self._encrypt_type = encrypt_type
         self._encryption: EncryptionStrategy | None = None
+        self._initial_session_key: bytes = b""
         self._simple_assembler = SimplePacketAssembler()
         self._options = Connection.Options()
 
@@ -631,6 +632,11 @@ def _set_state(
         if state.is_error:
             self._notify_disconnect(exc)
 
+    def set_state(
+        self, state: ConnectionState, exc: Exception | type[Exception] | None = None
+    ) -> None:
+        self._set_state(state, exc)
+
     def _get_characteristics(self, char_type: Literal["write", "notify"]):
         assert self._client is not None
 
@@ -944,6 +950,7 @@ async def getKeyInfoReqHandler(
 
         # Parse the data that contains sRand (first 16 bytes) & seed (last 2 bytes)
         session_key = await self.genSessionKey(data[16:18], data[:16])
+        self._initial_session_key = self._encryption.session_key
         self._encryption = Type7Encryption(session_key, self._encryption.iv)
 
         await self.getAuthStatus()

custom_components/ef_ble/eflib/devicebase.py

@@ -139,6 +139,15 @@ def auth_header_dst(self) -> int:
     def connection_state(self):
         return None if self._conn is None else self._conn._connection_state
 
+    def set_connection_state(
+        self,
+        state: ConnectionState,
+        exc: Exception | type[Exception] | None = None,
+    ) -> None:
+        if self._conn is None:
+            return
+        self._conn.set_state(state, exc)
+
     @property
     def diagnostics(self):
         return self._diagnostics
@@ -220,6 +229,11 @@ def with_enabled_packet_diagnostics(
         self._diagnostics.with_buffer_size(buffer_size)
         return self
 
+    def with_diagnostics_on_exception(self, enabled: bool = True):
+        """Enable automatic diagnostics save to disk on connection errors"""
+        self._diagnostics.with_save_on_exception(enabled)
+        return self
+
     def with_name(self, name: str):
         self._name = name
         return self

custom_components/ef_ble/eflib/logging_util.py

@@ -1,11 +1,14 @@
+import asyncio
 import dataclasses
 import json
 import logging
 import re
 import time
+import traceback
 from collections import deque
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from datetime import UTC, datetime
 from enum import Flag, auto
 from functools import cached_property
 from pathlib import Path
@@ -268,8 +271,10 @@ class DeviceDiagnostics:
     disconnect_times: list[float]
     raw_data_connection: list[tuple[float, bytes]]
     raw_data_messages: list[tuple[float, bytes]]
+    raw_data_send: list[tuple[float, bytes]]
     iv: bytes
     session_key: bytes
+    initial_session_key: bytes
 
     def _encode_bytes(self, value: bytes, session: Session | None) -> str:
         if session is not None:
@@ -289,8 +294,12 @@ def serialize(self, session: Session | None = None):
             raw_data_messages=[
                 (k, self._encode_bytes(v, session)) for (k, v) in self.raw_data_messages
             ],
+            raw_data_send=[
+                (k, self._encode_bytes(v, session)) for (k, v) in self.raw_data_send
+            ],
             iv=self._encode_bytes(self.iv, session),
             session_key=self._encode_bytes(self.session_key, session),
+            initial_session_key=self._encode_bytes(self.initial_session_key, session),
         )
 
     def as_dict(self):
@@ -304,36 +313,66 @@ class DeviceDiagnosticsCollector:
     def __init__(self, device: "DeviceBase", buffer_size: int = 100):
         self._device = device
         self._enabled = False
+        self._save_on_exception = False
         self._buffer_size = buffer_size
 
         self._last_packets: deque[tuple[float, bytes]] = deque(maxlen=buffer_size)
         self._last_errors: deque[tuple[float, str]] = deque(maxlen=buffer_size)
         self._connect_times: deque[float] = deque(maxlen=buffer_size)
         self._raw_data_connection: list[tuple[float, bytes]] = []
         self._raw_data_messages: deque[tuple[float, bytes]] = deque(maxlen=1000)
+        self._raw_data_send: deque[tuple[float, bytes]] = deque(maxlen=1000)
 
         self._disconnect_times: deque[float] = deque(maxlen=buffer_size)
         self._skip_first_messages: int = 8
         self._unlisten_callbacks: list[Callable[[], None]] = []
 
         self._start_time = time.time()
 
+        self._logger = logging.getLogger(__name__)
+
     def as_dict(self, session: Session | None = None):
         """Get diagnostics data as dictionary"""
         return self.diagnostics.serialize(session).as_dict()
 
+    def build_diagnostics_dict(self, session: Session | None = None) -> dict:
+        device = self._device
+        result: dict = {
+            "device": device.device,
+            "name": device.name,
+            "default_name": device._default_name,
+            "sn_prefix": device._sn[:4],
+            "connection_state": device.connection_state,
+            "connection_state_history": list(device.connection_log.history),
+            "manufacturer_data": (
+                session.encrypt(device._manufacturer_data).hex()
+                if session is not None
+                else device._manufacturer_data.hex()
+            ),
+        }
+        if session is not None:
+            result["session"] = session.header.hex()
+        if self.is_enabled:
+            result |= self.as_dict(session)
+        return result
+
     @property
     def diagnostics(self):
         """Get diagnostics data"""
+        conn = self._device._conn
+        encryption = conn._encryption
+
         return DeviceDiagnostics(
             last_packets=list(self._last_packets),
             last_errors=list(self._last_errors),
             connect_times=list(self._connect_times),
             disconnect_times=list(self._disconnect_times),
             raw_data_connection=self._raw_data_connection,
             raw_data_messages=list(self._raw_data_messages),
-            iv=self._device._conn._encryption.iv,
-            session_key=self._device._conn._encryption.session_key,
+            raw_data_send=list(self._raw_data_send),
+            iv=encryption.iv if encryption is not None else b"",
+            session_key=encryption.session_key if encryption is not None else b"",
+            initial_session_key=conn._initial_session_key,
         )
 
     @property
@@ -368,6 +407,7 @@ def enabled(self, enabled: bool = True):
                     self._device.on_packet_received(self._on_packet_received),
                     self._device.on_packet_parsed(self._on_packet_parsed),
                     self._device.on_data_received(self._on_data_received),
+                    self._device.on_data_send(self._on_data_send),
                 ]
             )
             return self
@@ -448,11 +488,97 @@ def _on_data_received(self, data: bytes, state: "ConnectionState"):
 
         buffer.append(self._with_time(data))
 
+    def _on_data_send(self, data: bytes):
+        self._raw_data_send.append(self._with_time(data))
+
+    def with_save_on_exception(self, enabled: bool = True):
+        """
+        Enable or disable automatic diagnostics save on connection error
+
+        When enabled, packet collection is force-enabled and a state change
+        listener is registered. On any error state, collected diagnostics
+        are saved to disk automatically.
+        """
+        if enabled == self._save_on_exception:
+            return self
+
+        self._save_on_exception = enabled
+
+        if enabled:
+            self.enabled(True)
+            self._unlisten_callbacks.append(
+                self._device.on_connection_state_change(self._on_state_change)
+            )
+
+        return self
+
+    def _on_state_change(self, state: "ConnectionState") -> None:
+        if not self._save_on_exception or not state.is_error:
+            return
+
+        conn = self._device._conn
+        exc = getattr(conn, "_last_exception", None) if conn is not None else None
+        try:
+            self._save_to_disk(state, exc)
+        except Exception:
+            self._device._logger.exception(
+                "Failed to save diagnostics-on-exception snapshot"
+            )
+
+    def _save_to_disk(
+        self,
+        state: "ConnectionState",
+        exc: Exception | type[Exception] | None = None,
+    ) -> None:
+        session = Session()
+
+        exc_message = None
+        exc_traceback = None
+        if exc is not None:
+            exc_message = session.encrypt(str(exc).encode()).hex()
+        if isinstance(exc, BaseException) and exc.__traceback__ is not None:
+            tb = "".join(traceback.format_tb(exc.__traceback__))
+            exc_traceback = session.encrypt(tb.encode()).hex()
+
+        data = {
+            "timestamp": datetime.now(UTC).isoformat(),
+            "exception": {
+                "type": type(exc).__name__ if exc is not None else None,
+                "message": exc_message,
+                "traceback": exc_traceback,
+                "state_on_error": state.name,
+            },
+            "data": self.build_diagnostics_dict(session),
+        }
+
+        sn_prefix = self._device._sn[:4]
+        ts = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
+        cache_dir = Path(__file__).parent.parent / ".diagnostics"
+        path = cache_dir / f"{sn_prefix}_exception_{ts}.json"
+        content = json.dumps(data, default=str, indent=2)
+
+        def _write() -> None:
+            cache_dir.mkdir(exist_ok=True)
+            path.write_text(content)
+
+        task = asyncio.get_running_loop().run_in_executor(None, _write)
+
+        def _log_result(future: asyncio.Future) -> None:
+            if (err := future.exception()) is not None:
+                self._device._logger.error(
+                    "Failed to save diagnostics to %s: %s", path, err
+                )
+            else:
+                self._device._logger.info("Diagnostics saved to %s", path)
+
+        task.add_done_callback(_log_result)
+
     def _clear_buffers(self):
         self._last_packets.clear()
         self._last_errors.clear()
         self._connect_times.clear()
         self._disconnect_times.clear()
+        self._raw_data_send.clear()
 
 
 class _LazyHex: