Simplify StoredDevice and serialize backup JSON only at the very end

Author: puddly

zigpy_znp/api.py

@@ -6,6 +6,7 @@
 import logging
 import itertools
 import contextlib
+import dataclasses
 from collections import Counter, defaultdict
 
 import zigpy.state
@@ -165,28 +166,14 @@ async def load_network_info(self, *, load_devices=False):
 
         # This takes a few seconds
         if load_devices:
-            for device in await security.read_devices(self):
-                if device.is_child:
-                    network_info.children.append(
-                        zigpy.state.NodeInfo(
-                            ieee=device.ieee,
-                            nwk=device.nwk,
-                            logical_type=zdo_t.LogicalType.EndDevice,
-                        )
-                    )
+            for dev in await security.read_devices(self):
+                if dev.is_child:
+                    network_info.children.append(dev.node_info)
                 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,
-                        )
-                    )
+                    network_info.nwk_addresses[dev.node_info.ieee] = dev.node_info.nwk
+
+                if dev.key is not None:
+                    network_info.key_table.append(dev.key)
 
         ieee = await self.nvram.osal_read(OsalNvIds.EXTADDR, item_type=t.EUI64)
         logical_type = await self.nvram.osal_read(
@@ -395,8 +382,10 @@ async def write_network_info(
 
         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,
+                node_info=dataclasses.replace(
+                    child, nwk=0xFFFE if child.nwk is None else child.nwk
+                ),
+                key=None,
                 is_child=True,
             )
 
@@ -405,26 +394,30 @@ async def write_network_info(
 
             if device is None:
                 device = security.StoredDevice(
-                    nwk=network_info.nwk_addresses.get(key.partner_ieee, 0xFFFE),
-                    ieee=key.partner_ieee,
+                    node_info=zigpy.state.NodeInfo(
+                        nwk=network_info.nwk_addresses.get(key.partner_ieee, 0xFFFE),
+                        ieee=key.partner_ieee,
+                        logical_type=None,
+                    ),
+                    key=None,
                     is_child=False,
                 )
 
-            devices[key.partner_ieee] = device.replace(
-                aps_link_key=key.key,
-                tx_counter=key.tx_counter,
-                rx_counter=key.rx_counter,
-            )
+            devices[key.partner_ieee] = device.replace(key=key)
 
         LOGGER.debug("Writing children and keys")
 
-        await security.write_devices(
+        new_tclk_seed = await security.write_devices(
             znp=self,
             devices=list(devices.values()),
             tclk_seed=tclk_seed,
             counter_increment=0,
         )
 
+        # If the provided TCLK seed isn't optimal, overwrite it
+        if new_tclk_seed != tclk_seed:
+            await self.nvram.osal_write(OsalNvIds.TCLK_SEED, new_tclk_seed)
+
         if self.version == 1.2:
             await self.nvram.osal_write(
                 OsalNvIds.HAS_CONFIGURED_ZSTACK1,

zigpy_znp/tools/network_backup.py

@@ -6,76 +6,100 @@
 import logging
 import datetime
 
+import zigpy.state
+
 import zigpy_znp
 import zigpy_znp.types as t
 from zigpy_znp.api import ZNP
 from zigpy_znp.tools.common import ClosableFileType, setup_parser, validate_backup_json
-from zigpy_znp.znp.security import read_devices
 from zigpy_znp.zigbee.application import ControllerApplication
 
 LOGGER = logging.getLogger(__name__)
 
 
-async def backup_network(znp: ZNP) -> t.JSONType:
-    try:
-        await znp.load_network_info()
-    except ValueError as e:
-        raise RuntimeError("Failed to load network info") from e
-
-    devices = []
+def zigpy_state_to_json_backup(
+    network_info: zigpy.state.NetworkInformation, node_info: zigpy.state.NodeInfo
+) -> t.JSONType:
+    devices = {}
 
-    for device in await read_devices(znp):
-        nwk = device.nwk if device.nwk != 0xFFFE else None
-        obj = {
+    for ieee, nwk in network_info.nwk_addresses.items():
+        devices[ieee] = {
+            "ieee_address": ieee.serialize()[::-1].hex(),
             "nwk_address": nwk.serialize()[::-1].hex(),
-            "ieee_address": device.ieee.serialize()[::-1].hex(),
-            "is_child": device.is_child,
+            "is_child": False,
         }
 
-        if device.aps_link_key:
-            obj["link_key"] = {
-                "tx_counter": device.tx_counter,
-                "rx_counter": device.rx_counter,
-                "key": device.aps_link_key.serialize().hex(),
-            }
-
-        devices.append(obj)
+    for child in network_info.children:
+        devices[child.ieee] = {
+            "ieee_address": child.ieee.serialize()[::-1].hex(),
+            "nwk_address": child.nwk.serialize()[::-1].hex()
+            if child.nwk is not None
+            else None,
+            "is_child": True,
+        }
 
-    devices.sort(key=lambda d: d["ieee_address"])
+    for key in network_info.key_table:
+        if key.partner_ieee not in devices:
+            devices[key.partner_ieee] = {
+                "ieee_address": key.partner_ieee.serialize()[::-1].hex(),
+                "nwk_address": None,
+                "is_child": False,
+            }
 
-    now = datetime.datetime.now().astimezone()
+        devices[key.partner_ieee]["link_key"] = {
+            "key": key.key.serialize().hex(),
+            "tx_counter": key.tx_counter,
+            "rx_counter": key.rx_counter,
+        }
 
-    obj = {
+    return {
         "metadata": {
             "version": 1,
             "format": "zigpy/open-coordinator-backup",
-            "source": f"zigpy-znp@{zigpy_znp.__version__}",
-            "internal": {
-                "creation_time": now.isoformat(timespec="seconds"),
-                "zstack": {
-                    "version": znp.version,
-                },
-            },
+            "source": None,
+            "internal": None,
         },
-        "coordinator_ieee": znp.node_info.ieee.serialize()[::-1].hex(),
-        "pan_id": znp.network_info.pan_id.serialize()[::-1].hex(),
-        "extended_pan_id": znp.network_info.extended_pan_id.serialize()[::-1].hex(),
-        "nwk_update_id": znp.network_info.nwk_update_id,
-        "security_level": znp.network_info.security_level,
-        "channel": znp.network_info.channel,
-        "channel_mask": list(znp.network_info.channel_mask),
+        "coordinator_ieee": node_info.ieee.serialize()[::-1].hex(),
+        "pan_id": network_info.pan_id.serialize()[::-1].hex(),
+        "extended_pan_id": network_info.extended_pan_id.serialize()[::-1].hex(),
+        "nwk_update_id": network_info.nwk_update_id,
+        "security_level": network_info.security_level,
+        "channel": network_info.channel,
+        "channel_mask": list(network_info.channel_mask),
         "network_key": {
-            "key": znp.network_info.network_key.key.serialize().hex(),
-            "sequence_number": znp.network_info.network_key.seq,
-            "frame_counter": znp.network_info.network_key.tx_counter,
+            "key": network_info.network_key.key.serialize().hex(),
+            "sequence_number": network_info.network_key.seq,
+            "frame_counter": network_info.network_key.tx_counter,
+        },
+        "devices": sorted(devices.values(), key=lambda d: d["ieee_address"]),
+    }
+
+
+async def backup_network(znp: ZNP) -> t.JSONType:
+    try:
+        await znp.load_network_info()
+    except ValueError as e:
+        raise RuntimeError("Failed to load network info") from e
+
+    await znp.load_network_info(load_devices=True)
+
+    obj = zigpy_state_to_json_backup(
+        network_info=znp.network_info,
+        node_info=znp.node_info,
+    )
+
+    now = datetime.datetime.now().astimezone()
+
+    obj["metadata"]["source"] = f"zigpy-znp@{zigpy_znp.__version__}"
+    obj["metadata"]["internal"] = {
+        "creation_time": now.isoformat(timespec="seconds"),
+        "zstack": {
+            "version": znp.version,
         },
-        "devices": devices,
     }
 
-    if znp.network_info.stack_specific.get("zstack", {}).get("tclk_seed"):
-        obj.setdefault("stack_specific", {}).setdefault("zstack", {})[
-            "tclk_seed"
-        ] = znp.network_info.stack_specific["zstack"]["tclk_seed"]
+    if znp.network_info.stack_specific:
+        obj["stack_specific"] = znp.network_info.stack_specific
 
     # Ensure our generated backup is valid
     validate_backup_json(obj)

zigpy_znp/znp/security.py

@@ -16,14 +16,8 @@
 
 @dataclasses.dataclass(frozen=True)
 class StoredDevice:
-    ieee: t.EUI64
-    nwk: t.NWK
-
-    aps_link_key: t.KeyData = None
-
-    tx_counter: t.uint32_t = None
-    rx_counter: t.uint32_t = None
-
+    node_info: zigpy.state.NodeInfo
+    key: zigpy.state.Key | None
     is_child: bool = False
 
     def replace(self, **kwargs) -> StoredDevice:
@@ -63,7 +57,7 @@ def iter_seed_candidates(
 ) -> typing.Iterable[tuple[int, t.KeyData]]:
     for ieee, key in ieees_and_keys:
         # Derive a seed from each candidate. All rotations of a seed are equivalent.
-        tclk_seed = compute_tclk_seed(ieee, key, 0)
+        tclk_seed = t.KeyData(compute_tclk_seed(ieee, key, 0))
 
         # And see how many other keys share this same seed
         count = count_seed_matches(ieees_and_keys, tclk_seed)
@@ -276,8 +270,12 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
             t.AddrMgrUserType.Security,
         ):
             devices[entry.extAddr] = StoredDevice(
-                ieee=entry.extAddr,
-                nwk=entry.nwkAddr,
+                node_info=zigpy.state.NodeInfo(
+                    nwk=entry.nwkAddr,
+                    ieee=entry.extAddr,
+                    logical_type=None,
+                ),
+                key=None,
                 is_child=bool(entry.type & t.AddrMgrUserType.Assoc),
             )
         else:
@@ -294,11 +292,7 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
             )
             continue
 
-        devices[key.partner_ieee] = devices[key.partner_ieee].replace(
-            tx_counter=key.tx_counter,
-            rx_counter=key.rx_counter,
-            aps_link_key=key.key,
-        )
+        devices[key.partner_ieee] = devices[key.partner_ieee].replace(key=key)
 
     async for key in read_unhashed_link_keys(znp, addr_mgr):
         if key.partner_ieee not in devices:
@@ -311,11 +305,7 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
             )
             continue
 
-        devices[key.partner_ieee] = devices[key.partner_ieee].replace(
-            tx_counter=key.tx_counter,
-            rx_counter=key.rx_counter,
-            aps_link_key=key.key,
-        )
+        devices[key.partner_ieee] = devices[key.partner_ieee].replace(key=key)
 
     return list(devices.values())
 
@@ -328,11 +318,11 @@ async def write_addr_manager_entries(
     for dev in devices:
         entry = t.AddrMgrEntry(
             type=t.AddrMgrUserType.Default,
-            nwkAddr=dev.nwk,
-            extAddr=dev.ieee,
+            nwkAddr=dev.node_info.nwk,
+            extAddr=dev.node_info.ieee,
         )
 
-        if dev.aps_link_key:
+        if dev.key is not None:
             entry.type |= t.AddrMgrUserType.Security
 
         if dev.is_child:
@@ -366,9 +356,9 @@ async def write_devices(
     znp: ZNP,
     devices: typing.Sequence[StoredDevice],
     counter_increment: t.uint32_t = 2500,
-    tclk_seed: bytes = None,
-) -> None:
-    ieees_and_keys = [(d.ieee, d.aps_link_key) for d in devices if d.aps_link_key]
+    tclk_seed: t.KeyData = None,
+) -> t.KeyData:
+    ieees_and_keys = [(d.node_info.ieee, d.key.key) for d in devices if d.key]
 
     # Find the tclk_seed that maximizes the number of keys that can be derived from it
     if ieees_and_keys:
@@ -397,19 +387,19 @@ async def write_devices(
     aps_key_data_table = []
     link_key_table = t.APSLinkKeyTable()
 
-    for index, device in enumerate(devices):
-        if not device.aps_link_key:
+    for index, dev in enumerate(devices):
+        if dev.key is None:
             continue
 
-        shift = find_key_shift(device.ieee, device.aps_link_key, tclk_seed)
+        shift = find_key_shift(dev.node_info.ieee, dev.key.key, tclk_seed)
 
         if shift is not None:
             # Hashed link keys can be written into the TCLK table
             hashed_link_key_table.append(
                 t.TCLKDevEntry(
-                    txFrmCntr=device.tx_counter + counter_increment,
-                    rxFrmCntr=device.rx_counter,
-                    extAddr=device.ieee,
+                    txFrmCntr=dev.key.tx_counter + counter_increment,
+                    rxFrmCntr=dev.key.rx_counter,
+                    extAddr=dev.node_info.ieee,
                     keyAttributes=t.KeyAttributes.VERIFIED_KEY,
                     keyType=t.KeyType.NONE,
                     SeedShift_IcIndex=shift,
@@ -419,9 +409,9 @@ async def write_devices(
             # Unhashed link keys are written to another table
             aps_key_data_table.append(
                 t.APSKeyDataTableEntry(
-                    Key=device.aps_link_key,
-                    TxFrameCounter=device.tx_counter + counter_increment,
-                    RxFrameCounter=device.rx_counter,
+                    Key=dev.key.key,
+                    TxFrameCounter=dev.key.tx_counter + counter_increment,
+                    RxFrameCounter=dev.key.rx_counter,
                 )
             )
 
@@ -463,7 +453,7 @@ async def write_devices(
     await write_addr_manager_entries(znp, devices)
 
     if old_link_key_table is None:
-        return
+        return tclk_seed
 
     await znp.nvram.osal_write(OsalNvIds.APS_LINK_KEY_TABLE, new_link_key_table_value)
 
@@ -508,3 +498,5 @@ async def write_devices(
             values=aps_key_data_table,
             fill_value=aps_key_data_fill_value,
         )
+
+    return tclk_seed