Incorporate latest zigpy.state changes

Author: puddly

zigpy_znp/api.py

@@ -82,7 +82,7 @@ async def detect_zstack_version(self) -> float:
         except CommandNotRecognized:
             return 3.0
 
-    async def load_network_info(self, *, load_keys=False):
+    async def load_network_info(self, *, load_devices=False):
         """
         Loads low-level network information from NVRAM.
         Loading key data greatly increases the runtime so it not enabled by default.
@@ -150,6 +150,8 @@ async def load_network_info(self, *, load_keys=False):
                 partner_ieee=None,
             ),
             tc_link_key=None,
+            children=[],
+            nwk_addresses={},
             key_table=[],
             stack_specific=None,
         )
@@ -162,17 +164,29 @@ async def load_network_info(self, *, load_keys=False):
             }
 
         # This takes a few seconds
-        if load_keys:
+        if load_devices:
             for device in await security.read_devices(self):
-                network_info.key_table.append(
-                    zigpy.state.Key(
-                        key=device.aps_link_key,
-                        seq=0,
-                        tx_counter=device.tx_counter,
-                        rx_counter=device.rx_counter,
-                        partner_ieee=device.ieee,
+                if device.is_child:
+                    network_info.children.append(
+                        zigpy.state.NodeInfo(
+                            ieee=device.ieee,
+                            nwk=device.nwk,
+                            logical_type=zdo_t.LogicalType.EndDevice,
+                        )
+                    )
+                else:
+                    network_info.nwk_addresses[device.ieee] = device.nwk
+
+                if device.aps_link_key is not None:
+                    network_info.key_table.append(
+                        zigpy.state.Key(
+                            key=device.aps_link_key,
+                            seq=0,
+                            tx_counter=device.tx_counter,
+                            rx_counter=device.rx_counter,
+                            partner_ieee=device.ieee,
+                        )
                     )
-                )
 
         ieee = await self.nvram.osal_read(OsalNvIds.EXTADDR, item_type=t.EUI64)
         logical_type = await self.nvram.osal_read(
@@ -379,30 +393,21 @@ async def write_network_info(
 
         devices = {}
 
-        for neighbor in network_info.neighbor_table or []:
-            devices[neighbor.ieee] = security.StoredDevice(
-                nwk=neighbor.nwk,
-                ieee=neighbor.ieee,
+        for child in network_info.children or []:
+            devices[child.ieee] = security.StoredDevice(
+                nwk=child.nwk if child.nwk is not None else 0xFFFE,
+                ieee=child.ieee,
                 is_child=True,
             )
 
         for key in network_info.key_table or []:
             device = devices.get(key.partner_ieee)
 
             if device is None:
-                nwk = None
-
-                # The key table doesn't save a device's NWK address so we pick one
-                for test_nwk in range(0x0001, 0xFFFE + 1):
-                    if all(n.nwk != test_nwk for n in network_info.neighbor_table):
-                        nwk = test_nwk
-                        break
-
-                assert nwk is not None
-
                 device = security.StoredDevice(
-                    nwk=nwk,
+                    nwk=network_info.nwk_addresses.get(key.partner_ieee, 0xFFFE),
                     ieee=key.partner_ieee,
+                    is_child=False,
                 )
 
             devices[key.partner_ieee] = device.replace(
@@ -411,7 +416,7 @@ async def write_network_info(
                 rx_counter=key.rx_counter,
             )
 
-        LOGGER.debug("Writing neighbors and keys")
+        LOGGER.debug("Writing children and keys")
 
         await security.write_devices(
             znp=self,

zigpy_znp/tools/common.py

@@ -68,7 +68,10 @@
             "items": {
                 "type": "object",
                 "properties": {
-                    "nwk_address": {"type": "string", "pattern": "[a-fA-F0-9]{4}"},
+                    "nwk_address": {
+                        "type": ["string", "null"],
+                        "pattern": "[a-fA-F0-9]{4}",
+                    },
                     "ieee_address": {"type": "string", "pattern": "[a-fA-F0-9]{16}"},
                     "is_child": {"type": "boolean"},
                     "link_key": {

zigpy_znp/tools/network_backup.py

@@ -25,12 +25,11 @@ async def backup_network(znp: ZNP) -> t.JSONType:
     devices = []
 
     for device in await read_devices(znp):
+        nwk = device.nwk if device.nwk != 0xFFFE else None
         obj = {
-            "nwk_address": device.nwk.serialize()[::-1].hex(),
+            "nwk_address": nwk.serialize()[::-1].hex(),
             "ieee_address": device.ieee.serialize()[::-1].hex(),
-            "is_child": any(
-                device.ieee == nbr.ieee for nbr in znp.network_info.neighbor_table
-            ),
+            "is_child": device.is_child,
         }
 
         if device.aps_link_key:

zigpy_znp/tools/network_restore.py

@@ -48,18 +48,20 @@ def json_backup_to_zigpy_state(
     network_info.network_key.partner_ieee = None
     network_info.network_key.seq = backup["network_key"]["sequence_number"]
 
-    network_info.key_table = []
-    network_info.neighbor_table = []
+    network_info.children = []
+    network_info.nwk_addresses = {}
 
     for obj in backup["devices"]:
         node = zigpy.state.NodeInfo()
         node.nwk, _ = t.NWK.deserialize(bytes.fromhex(obj["nwk_address"])[::-1])
         node.ieee, _ = t.EUI64.deserialize(bytes.fromhex(obj["ieee_address"])[::-1])
         node.logical_type = None
 
-        # The `is_child` key is optional. If it is missing, preserve the old behavior
-        if obj["is_child"] if "is_child" in obj else ("link_key" not in obj):
-            network_info.neighbor_table.append(node)
+        # The `is_child` key is currently optional
+        if obj.get("is_child", True):
+            network_info.children.append(node)
+        else:
+            network_info.nwk_addresses[node.ieee] = node.nwk
 
         if "link_key" in obj:
             key = zigpy.state.Key()

zigpy_znp/zigbee/application.py

@@ -386,8 +386,9 @@ async def form_network(self):
                 partner_ieee=None,
             ),
             tc_link_key=None,
+            children=[],
             key_table=[],
-            neighbor_table=[],
+            nwk_addresses={},
             stack_specific={"zstack": {"tclk_seed": os.urandom(16).hex()}},
         )
 

zigpy_znp/znp/security.py

@@ -275,10 +275,6 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
             t.AddrMgrUserType.Assoc | t.AddrMgrUserType.Security,
             t.AddrMgrUserType.Security,
         ):
-            if not 0x0000 <= entry.nwkAddr <= 0xFFF7:
-                LOGGER.warning("Ignoring invalid address manager entry: %s", entry)
-                continue
-
             devices[entry.extAddr] = StoredDevice(
                 ieee=entry.extAddr,
                 nwk=entry.nwkAddr,