Repair incorrect TCLK partner IEEE address on startup

bellows/ezsp/__init__.py

@@ -5,6 +5,7 @@
 import asyncio
 import collections
 import contextlib
+import dataclasses
 import functools
 import logging
 import sys
@@ -20,6 +21,7 @@
 
 import bellows.config as conf
 from bellows.exception import EzspError, InvalidCommandError
+from bellows.ezsp.config import DEFAULT_CONFIG, RuntimeConfig, ValueConfig
 import bellows.types as t
 import bellows.uart
 
@@ -133,15 +135,14 @@ async def probe(cls, device_config: dict) -> bool | dict[str, int | str | bool]:
     async def _probe(self) -> None:
         """Open port and try sending a command"""
         await self.connect(use_thread=False)
-        await self._startup_reset()
-        await self.version()
+        await self.startup_reset()
 
     @property
     def is_tcp_serial_port(self) -> bool:
         parsed_path = urllib.parse.urlparse(self._config[conf.CONF_DEVICE_PATH])
         return parsed_path.scheme == "socket"
 
-    async def _startup_reset(self):
+    async def startup_reset(self) -> None:
         """Start EZSP and reset the stack."""
         # `zigbeed` resets on startup
         if self.is_tcp_serial_port:
@@ -157,19 +158,16 @@ async def _startup_reset(self):
         if not self.is_ezsp_running:
             await self.reset()
 
+        await self.version()
+
     @classmethod
     async def initialize(cls, zigpy_config: dict) -> EZSP:
         """Return initialized EZSP instance."""
         ezsp = cls(zigpy_config[conf.CONF_DEVICE])
         await ezsp.connect(use_thread=zigpy_config[conf.CONF_USE_THREAD])
 
         try:
-            await ezsp._startup_reset()
-            await ezsp.version()
-            await ezsp._protocol.initialize(zigpy_config)
-
-            if zigpy_config[zigpy.config.CONF_SOURCE_ROUTING]:
-                await ezsp.set_source_routing()
+            await ezsp.startup_reset()
         except Exception:
             ezsp.close()
             raise
@@ -419,6 +417,20 @@ async def can_rewrite_custom_eui64(self) -> bool:
         """Checks if the device EUI64 can be written any number of times."""
         return await self._get_nv3_restored_eui64_key() is not None
 
+    async def reset_custom_eui64(self) -> None:
+        """Reset the custom EUI64, if possible."""
+
+        nv3_eui64_key = await self._get_nv3_restored_eui64_key()
+        if nv3_eui64_key is None:
+            return
+
+        (status,) = await self.setTokenData(
+            nv3_eui64_key,
+            0,
+            t.LVBytes32(t.EmberEUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF").serialize()),
+        )
+        assert status == t.EmberStatus.SUCCESS
+
     async def write_custom_eui64(
         self, ieee: t.EUI64, *, burn_into_userdata: bool = False
     ) -> None:
@@ -488,7 +500,9 @@ async def set_source_routing(self) -> None:
         LOGGER.debug("Set concentrator type: %s", res)
         if res[0] != self.types.EmberStatus.SUCCESS:
             LOGGER.warning("Couldn't set concentrator type %s: %s", True, res)
-        await self._protocol.set_source_routing()
+
+        if self._ezsp_version >= 8:
+            await self.setSourceRouteDiscoveryMode(1)
 
     def start_ezsp(self):
         """Mark EZSP as running."""
@@ -512,3 +526,104 @@ def ezsp_version(self):
     def types(self):
         """Return EZSP types for this specific version."""
         return self._protocol.types
+
+    async def write_config(self, config: dict) -> None:
+        """Initialize EmberZNet Stack."""
+        config = self._protocol.SCHEMAS[conf.CONF_EZSP_CONFIG](config)
+
+        # Not all config will be present in every EZSP version so only use valid keys
+        ezsp_config = {}
+        ezsp_values = {}
+
+        for cfg in DEFAULT_CONFIG[self._ezsp_version]:
+            if isinstance(cfg, RuntimeConfig):
+                ezsp_config[cfg.config_id.name] = dataclasses.replace(
+                    cfg, config_id=self.types.EzspConfigId[cfg.config_id.name]
+                )
+            elif isinstance(cfg, ValueConfig):
+                ezsp_values[cfg.value_id.name] = dataclasses.replace(
+                    cfg, value_id=self.types.EzspValueId[cfg.value_id.name]
+                )
+
+        # Override the defaults with user-specified values (or `None` for deletions)
+        for name, value in config.items():
+            if value is None:
+                ezsp_config.pop(name)
+                continue
+
+            ezsp_config[name] = RuntimeConfig(
+                config_id=self.types.EzspConfigId[name],
+                value=value,
+            )
+
+        # Make sure CONFIG_PACKET_BUFFER_COUNT is always set last
+        if self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name in ezsp_config:
+            ezsp_config = {
+                **ezsp_config,
+                self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name: ezsp_config[
+                    self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name
+                ],
+            }
+
+        # First, set the values
+        for cfg in ezsp_values.values():
+            # XXX: A read failure does not mean the value is not writeable!
+            status, current_value = await self.getValue(cfg.value_id)
+
+            if status == self.types.EmberStatus.SUCCESS:
+                current_value, _ = type(cfg.value).deserialize(current_value)
+            else:
+                current_value = None
+
+            LOGGER.debug(
+                "Setting value %s = %s (old value %s)",
+                cfg.value_id.name,
+                cfg.value,
+                current_value,
+            )
+
+            (status,) = await self.setValue(cfg.value_id, cfg.value.serialize())
+
+            if status != self.types.EmberStatus.SUCCESS:
+                LOGGER.debug(
+                    "Could not set value %s = %s: %s",
+                    cfg.value_id.name,
+                    cfg.value,
+                    status,
+                )
+                continue
+
+        # Finally, set the config
+        for cfg in ezsp_config.values():
+            (status, current_value) = await self.getConfigurationValue(cfg.config_id)
+
+            # Only grow some config entries, all others should be set
+            if (
+                status == self.types.EmberStatus.SUCCESS
+                and cfg.minimum
+                and current_value >= cfg.value
+            ):
+                LOGGER.debug(
+                    "Current config %s = %s exceeds the default of %s, skipping",
+                    cfg.config_id.name,
+                    current_value,
+                    cfg.value,
+                )
+                continue
+
+            LOGGER.debug(
+                "Setting config %s = %s (old value %s)",
+                cfg.config_id.name,
+                cfg.value,
+                current_value,
+            )
+
+            (status,) = await self.setConfigurationValue(cfg.config_id, cfg.value)
+            if status != self.types.EmberStatus.SUCCESS:
+                LOGGER.debug(
+                    "Could not set config %s = %s: %s",
+                    cfg.config_id,
+                    cfg.value,
+                    status,
+                )
+                continue

bellows/ezsp/protocol.py

@@ -1,18 +1,17 @@
 import abc
 import asyncio
 import binascii
-import dataclasses
 import functools
 import logging
 import sys
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Tuple
 
 if sys.version_info[:2] < (3, 11):
     from async_timeout import timeout as asyncio_timeout  # pragma: no cover
 else:
     from asyncio import timeout as asyncio_timeout  # pragma: no cover
 
-from bellows.config import CONF_EZSP_CONFIG, CONF_EZSP_POLICIES
+from bellows.config import CONF_EZSP_POLICIES
 from bellows.exception import InvalidCommandError
 from bellows.typing import GatewayType
 
@@ -55,120 +54,6 @@ def _ezsp_frame_tx(self, name: str) -> bytes:
     async def pre_permit(self, time_s: int) -> None:
         """Schedule task before allowing new joins."""
 
-    async def initialize(self, zigpy_config: Dict) -> None:
-        """Initialize EmberZNet Stack."""
-
-        # Prevent circular import
-        from bellows.ezsp.config import DEFAULT_CONFIG, RuntimeConfig, ValueConfig
-
-        # Not all config will be present in every EZSP version so only use valid keys
-        ezsp_config = {}
-        ezsp_values = {}
-
-        for cfg in DEFAULT_CONFIG[self.VERSION]:
-            if isinstance(cfg, RuntimeConfig):
-                ezsp_config[cfg.config_id.name] = dataclasses.replace(
-                    cfg, config_id=self.types.EzspConfigId[cfg.config_id.name]
-                )
-            elif isinstance(cfg, ValueConfig):
-                ezsp_values[cfg.value_id.name] = dataclasses.replace(
-                    cfg, value_id=self.types.EzspValueId[cfg.value_id.name]
-                )
-
-        # Override the defaults with user-specified values (or `None` for deletions)
-        for name, value in self.SCHEMAS[CONF_EZSP_CONFIG](
-            zigpy_config[CONF_EZSP_CONFIG]
-        ).items():
-            if value is None:
-                ezsp_config.pop(name)
-                continue
-
-            ezsp_config[name] = RuntimeConfig(
-                config_id=self.types.EzspConfigId[name],
-                value=value,
-            )
-
-        # Make sure CONFIG_PACKET_BUFFER_COUNT is always set last
-        if self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name in ezsp_config:
-            ezsp_config = {
-                **ezsp_config,
-                self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name: ezsp_config[
-                    self.types.EzspConfigId.CONFIG_PACKET_BUFFER_COUNT.name
-                ],
-            }
-
-        # First, set the values
-        for cfg in ezsp_values.values():
-            # XXX: A read failure does not mean the value is not writeable!
-            status, current_value = await self.getValue(cfg.value_id)
-
-            if status == self.types.EmberStatus.SUCCESS:
-                current_value, _ = type(cfg.value).deserialize(current_value)
-            else:
-                current_value = None
-
-            LOGGER.debug(
-                "Setting value %s = %s (old value %s)",
-                cfg.value_id.name,
-                cfg.value,
-                current_value,
-            )
-
-            (status,) = await self.setValue(cfg.value_id, cfg.value.serialize())
-
-            if status != self.types.EmberStatus.SUCCESS:
-                LOGGER.debug(
-                    "Could not set value %s = %s: %s",
-                    cfg.value_id.name,
-                    cfg.value,
-                    status,
-                )
-                continue
-
-        # Finally, set the config
-        for cfg in ezsp_config.values():
-            (status, current_value) = await self.getConfigurationValue(cfg.config_id)
-
-            # Only grow some config entries, all others should be set
-            if (
-                status == self.types.EmberStatus.SUCCESS
-                and cfg.minimum
-                and current_value >= cfg.value
-            ):
-                LOGGER.debug(
-                    "Current config %s = %s exceeds the default of %s, skipping",
-                    cfg.config_id.name,
-                    current_value,
-                    cfg.value,
-                )
-                continue
-
-            LOGGER.debug(
-                "Setting config %s = %s (old value %s)",
-                cfg.config_id.name,
-                cfg.value,
-                current_value,
-            )
-
-            (status,) = await self.setConfigurationValue(cfg.config_id, cfg.value)
-            if status != self.types.EmberStatus.SUCCESS:
-                LOGGER.debug(
-                    "Could not set config %s = %s: %s",
-                    cfg.config_id,
-                    cfg.value,
-                    status,
-                )
-                continue
-
-    async def get_free_buffers(self) -> Optional[int]:
-        status, value = await self.getValue(self.types.EzspValueId.VALUE_FREE_BUFFERS)
-
-        if status != self.types.EzspStatus.SUCCESS:
-            LOGGER.debug("Couldn't get free buffers: %s", status)
-            return None
-
-        return int.from_bytes(value, byteorder="little")
-
     async def command(self, name, *args) -> Any:
         """Serialize command and send it."""
         LOGGER.debug("Send command %s: %s", name, args)
@@ -182,13 +67,10 @@ async def command(self, name, *args) -> Any:
         async with asyncio_timeout(EZSP_CMD_TIMEOUT):
             return await future
 
-    async def set_source_routing(self) -> None:
-        """Enable source routing on NCP."""
-
-    async def update_policies(self, zigpy_config: dict) -> None:
+    async def update_policies(self, policy_config: dict) -> None:
         """Set up the policies for what the NCP should do."""
 
-        policies = self.SCHEMAS[CONF_EZSP_POLICIES](zigpy_config[CONF_EZSP_POLICIES])
+        policies = self.SCHEMAS[CONF_EZSP_POLICIES](policy_config)
         self.tc_policy = policies[self.types.EzspPolicyId.TRUST_CENTER_POLICY.name]
 
         for policy, value in policies.items():