Enable ALLOW_UNSECURED_REJOINS for EZSPv8+

bellows/ezsp/__init__.py

@@ -5,7 +5,7 @@
 import asyncio
 from asyncio import timeout as asyncio_timeout
 import collections
-from collections.abc import Callable, Generator
+from collections.abc import AsyncIterator, Callable, Generator
 import contextlib
 import dataclasses
 import functools
@@ -18,6 +18,7 @@
 from bellows.ash import NcpFailure
 import bellows.config as conf
 from bellows.exception import (
+    ControllerError,
     EzspError,
     InvalidCommandError,
     InvalidCommandPayload,
@@ -52,6 +53,9 @@
 NETWORK_OPS_TIMEOUT = 10
 NETWORK_COORDINATOR_STARTUP_RESET_WAIT = 2
 
+WILDCARD_EUI64 = t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF")
+WELL_KNOWN_LINK_KEY = t.KeyData(b"ZigBeeAlliance09")
+
 
 class EZSP:
     _BY_VERSION = {
@@ -628,6 +632,37 @@ async def set_source_routing(self, enabled: bool) -> None:
                 )
             )
 
+    @contextlib.asynccontextmanager
+    async def allow_joins(
+        self, *, link_keys: dict[t.EUI64, t.KeyData] | None = None
+    ) -> AsyncIterator[None]:
+        """Allow joins with the well-known key, or only with the given link keys."""
+        install_codes_only = link_keys is not None
+
+        if link_keys is None:
+            link_keys = {WILDCARD_EUI64: WELL_KNOWN_LINK_KEY}
+
+        try:
+            # Keys from a previous window may remain if its cleanup failed
+            await self._protocol.clear_transient_link_keys()
+
+            for ieee, key in link_keys.items():
+                status = await self._protocol.add_transient_link_key(ieee, key)
+
+                if status != t.sl_Status.OK:
+                    raise ControllerError(f"Failed to set link key for {ieee}")
+
+            await self._protocol.set_join_policy(
+                allow_joins=True, install_codes_only=install_codes_only
+            )
+
+            yield
+        finally:
+            try:
+                await self._protocol.clear_transient_link_keys()
+            finally:
+                await self._protocol.set_join_policy(allow_joins=False)
+
     def start_ezsp(self):
         """Mark EZSP as running."""
         self._ezsp_event.set()

bellows/ezsp/protocol.py

@@ -283,15 +283,22 @@ def __getattr__(self, name: str) -> Callable:
 
         return functools.partial(self.command, name)
 
-    async def pre_permit(self, time_s: int) -> None:
-        """Schedule task before allowing new joins."""
+    @abc.abstractmethod
+    async def set_join_policy(
+        self, *, allow_joins: bool, install_codes_only: bool = False
+    ) -> None:
+        raise NotImplementedError
 
     @abc.abstractmethod
     async def add_transient_link_key(
         self, ieee: t.EUI64, key: t.KeyData
     ) -> t.sl_Status:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def clear_transient_link_keys(self) -> None:
+        raise NotImplementedError
+
     @abc.abstractmethod
     async def read_child_data(
         self,

bellows/ezsp/v10/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v9 import config as v9_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v11/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v10 import config as v10_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v12/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v11 import config as v11_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v13/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v12 import config as v12_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v14/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v13 import config as v13_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v16/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v14 import config as v14_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v17/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v16 import config as v16_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v18/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v17 import config as v17_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v19/config.py

@@ -3,7 +3,7 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v8.config import EZSP_POLICIES_SHARED
 from ..v18 import config as v18_config
 
 EZSP_SCHEMA = {

bellows/ezsp/v4/__init__.py

@@ -36,14 +36,19 @@ def _ezsp_frame_rx(self, data: bytes) -> tuple[int, int, bytes]:
         """Handler for received data frame."""
         return data[0], data[2], data[3:]
 
-    async def pre_permit(self, time_s: int) -> None:
+    async def set_join_policy(
+        self, *, allow_joins: bool, install_codes_only: bool = False
+    ) -> None:
         pass
 
     async def add_transient_link_key(
         self, ieee: t.EUI64, key: t.KeyData
     ) -> t.sl_Status:
         return t.sl_Status.OK
 
+    async def clear_transient_link_keys(self) -> None:
+        pass
+
     async def read_child_data(
         self,
     ) -> AsyncGenerator[tuple[t.NWK, t.EUI64, t.EmberNodeType]]:

bellows/ezsp/v5/__init__.py

@@ -41,11 +41,8 @@ async def add_transient_link_key(
         (status,) = await self.addTransientLinkKey(partner=ieee, transientKey=key)
         return t.sl_Status.from_ember_status(status)
 
-    async def pre_permit(self, time_s: int) -> None:
-        """Add pre-shared TC Link key."""
-        wild_card_ieee = t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF")
-        tc_link_key = t.KeyData(b"ZigBeeAlliance09")
-        await self.add_transient_link_key(wild_card_ieee, tc_link_key)
+    async def clear_transient_link_keys(self) -> None:
+        await self.clearTransientLinkKeys()
 
     async def read_address_table(self) -> AsyncGenerator[tuple[t.NWK, t.EUI64]]:
         (status, addr_table_size) = await self.getConfigurationValue(

bellows/ezsp/v8/__init__.py

@@ -1,5 +1,4 @@
 """"EZSP Protocol version 8 protocol handler."""
-import asyncio
 import logging
 
 import voluptuous
@@ -36,18 +35,25 @@ def _ezsp_frame_rx(self, data: bytes) -> tuple[int, int, bytes]:
 
         return seq, frame_id, data
 
-    async def pre_permit(self, time_s: int) -> None:
-        """Temporarily change TC policy while allowing new joins."""
-        await super().pre_permit(time_s)
-        await self.setPolicy(
-            policyId=t.EzspPolicyId.TRUST_CENTER_POLICY,
-            decisionId=(
-                t.EzspDecisionBitmask.ALLOW_JOINS
+    async def set_join_policy(
+        self, *, allow_joins: bool, install_codes_only: bool = False
+    ) -> None:
+        if not allow_joins:
+            decision = self.tc_policy
+        elif install_codes_only:
+            decision = (
+                self.tc_policy
+                | t.EzspDecisionBitmask.ALLOW_JOINS
                 | t.EzspDecisionBitmask.ALLOW_UNSECURED_REJOINS
-            ),
-        )
-        await asyncio.sleep(time_s + 2)
+                | t.EzspDecisionBitmask.JOINS_USE_INSTALL_CODE_KEY
+            )
+        else:
+            decision = (
+                self.tc_policy
+                | t.EzspDecisionBitmask.ALLOW_JOINS
+                | t.EzspDecisionBitmask.ALLOW_UNSECURED_REJOINS
+            )
+
         await self.setPolicy(
-            policyId=t.EzspPolicyId.TRUST_CENTER_POLICY,
-            decisionId=self.tc_policy,
+            policyId=t.EzspPolicyId.TRUST_CENTER_POLICY, decisionId=decision
         )

bellows/ezsp/v8/config.py

@@ -1,11 +1,30 @@
 import voluptuous as vol
 
 from bellows.config import cv_optional_int, cv_uint16, extend_vol_schema
-from bellows.types import EzspConfigId, EzspPolicyId
+from bellows.types import EzspConfigId, EzspDecisionBitmask, EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
+from ..v4 import config as v4_config
 from ..v7 import config as v7_config
 
+# EZSPv8 replaces the trust center decisions with a decision bitmask
+EZSP_POLICIES_SHARED = {
+    **{
+        key: value
+        for key, value in v4_config.EZSP_POLICIES_SHARED.items()
+        if key != EzspPolicyId.TRUST_CENTER_POLICY.name
+    },
+    # Same as the NCP's boot default and the SDK host framework. Trust Center rejoins
+    # are always accepted. New joins still need permit joining and the well-known
+    # transient key, which only exists while joins are allowed.
+    vol.Optional(
+        EzspPolicyId.TRUST_CENTER_POLICY.name,
+        default=(
+            EzspDecisionBitmask.ALLOW_JOINS
+            | EzspDecisionBitmask.ALLOW_UNSECURED_REJOINS
+        ),
+    ): cv_uint16,
+}
+
 EZSP_SCHEMA = extend_vol_schema(
     v7_config.EZSP_SCHEMA,
     {

bellows/ezsp/v9/config.py

@@ -3,8 +3,8 @@
 from bellows.config import cv_uint16
 from bellows.types import EzspPolicyId
 
-from ..v4.config import EZSP_POLICIES_SHARED
 from ..v8 import config as v8_config
+from ..v8.config import EZSP_POLICIES_SHARED
 
 EZSP_SCHEMA = {
     **v8_config.EZSP_SCHEMA,