Include a can_write_custom_eui64 stack-specific key in the backup (#…

…469)

* Include a `can_write_custom_eui64` stack-specific key in the backup

* Fix failing unit tests

bellows/ezsp/__init__.py

@@ -291,6 +291,12 @@ async def get_board_info(self) -> Tuple[str, str, str]:
             version = "unknown stack version"
         return tokens[0], tokens[1], version
 
+    async def can_write_custom_eui64(self) -> bool:
+        """Checks if the write-once custom EUI64 token has been written."""
+        (custom_eui_64,) = await self.getMfgToken(t.EzspMfgTokenId.MFG_CUSTOM_EUI_64)
+
+        return custom_eui_64 == b"\xFF" * 8
+
     def add_callback(self, cb):
         id_ = hash(cb)
         while id_ in self._callbacks:

bellows/zigbee/application.py

@@ -263,6 +263,7 @@ async def load_network_info(self, *, load_devices=False) -> None:
             tc_link_key.partner_ieee = self.state.node_info.ieee
 
         brd_manuf, brd_name, version = await self._get_board_info()
+        can_write_custom_eui64 = await ezsp.can_write_custom_eui64()
 
         self.state.network_info = zigpy.state.NetworkInfo(
             source=f"bellows@{bellows.__version__}",
@@ -285,6 +286,7 @@ async def load_network_info(self, *, load_devices=False) -> None:
                     "board": brd_name,
                     "version": version,
                     "stack_version": ezsp.ezsp_version,
+                    "can_write_custom_eui64": can_write_custom_eui64,
                 }
             },
         )
@@ -351,6 +353,7 @@ async def write_network_info(
         except bellows.exception.EzspError:
             pass
 
+        can_write_custom_eui64 = await ezsp.can_write_custom_eui64()
         stack_specific = network_info.stack_specific.get("ezsp", {})
         (current_eui64,) = await ezsp.getEui64()
 
@@ -362,7 +365,12 @@ async def write_network_info(
                 "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
             )
 
-            if should_update_eui64:
+            if should_update_eui64 and not can_write_custom_eui64:
+                LOGGER.error(
+                    "Current node's IEEE address has already been written once. It"
+                    " cannot be written again without fully erasing the chip with JTAG."
+                )
+            elif should_update_eui64:
                 new_ncp_eui64 = t.EmberEUI64(node_info.ieee)
                 (status,) = await ezsp.setMfgToken(
                     t.EzspMfgTokenId.MFG_CUSTOM_EUI_64, new_ncp_eui64.serialize()

tests/test_application.py

@@ -126,6 +126,7 @@ async def mock_leave(*args, **kwargs):
     ezsp_mock.setConfigurationValue = AsyncMock(return_value=t.EmberStatus.SUCCESS)
     ezsp_mock.networkInit = AsyncMock(return_value=[init])
     ezsp_mock.getNetworkParameters = AsyncMock(return_value=[0, nwk_type, nwk_params])
+    ezsp_mock.can_write_custom_eui64 = AsyncMock(return_value=True)
 
     if board_info:
         ezsp_mock.get_board_info = AsyncMock(

tests/test_application_network_state.py

@@ -1,3 +1,5 @@
+import logging
+
 import pytest
 import zigpy.state
 import zigpy.types as zigpy_t
@@ -79,6 +81,7 @@ def _mock_app_for_load(app):
     ezsp = app._ezsp
 
     app._ensure_network_running = AsyncMock()
+    ezsp.can_write_custom_eui64 = AsyncMock(return_value=True)
     ezsp.getNetworkParameters = AsyncMock(
         return_value=[
             t.EmberStatus.SUCCESS,
@@ -364,6 +367,7 @@ def _mock_app_for_write(app, network_info, node_info, ezsp_ver=None):
     ezsp.formNetwork = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
     ezsp.setValue = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
     ezsp.setMfgToken = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+    ezsp.can_write_custom_eui64 = AsyncMock(return_value=True)
 
 
 async def test_write_network_info_failed_leave1(app, network_info, node_info):
@@ -415,6 +419,34 @@ async def test_write_network_info_write_new_eui64(app, network_info, node_info):
     )
 
 
+async def test_write_network_info_write_new_eui64_failure(
+    caplog, app, network_info, node_info
+):
+    _mock_app_for_write(app, network_info, node_info)
+
+    app._ezsp.can_write_custom_eui64.return_value = False
+
+    # Differs from what is in `node_info`
+    app._ezsp.getEui64.return_value = [t.EmberEUI64.convert("AA:AA:AA:AA:AA:AA:AA:AA")]
+
+    await app.write_network_info(
+        network_info=network_info.replace(
+            stack_specific={
+                "ezsp": {
+                    "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it": True,
+                    **network_info.stack_specific["ezsp"],
+                }
+            }
+        ),
+        node_info=node_info,
+    )
+
+    assert "cannot be written" in caplog.text
+
+    # The EUI64 is not written
+    app._ezsp.setMfgToken.assert_not_called()
+
+
 async def test_write_network_info_dont_write_new_eui64(app, network_info, node_info):
     _mock_app_for_write(app, network_info, node_info)
 

tests/test_ezsp.py

@@ -436,3 +436,16 @@ async def _mock_cmd(*args, **kwargs):
         cmd_mock.side_effect = _mock_cmd
         (status,) = await ezsp_f.leaveNetwork(timeout=0.01)
         assert status == t.EmberStatus.NETWORK_DOWN
+
+
+@pytest.mark.parametrize(
+    "value, expected_result",
+    [(b"\xFF" * 8, True), (bytes.fromhex("0846b8a11c004b1200"), False)],
+)
+async def test_can_write_custom_eui64(ezsp_f, value, expected_result):
+    ezsp_f.getMfgToken = AsyncMock(return_value=[value])
+
+    result = await ezsp_f.can_write_custom_eui64()
+    assert result == expected_result
+
+    ezsp_f.getMfgToken.assert_called_once_with(t.EzspMfgTokenId.MFG_CUSTOM_EUI_64)