Clean up TCLK seed recomputation

Author: puddly

tests/tools/test_network_backup_restore.py

@@ -303,26 +303,29 @@ async def test_tc_frame_counter_zstack33(make_connected_znp):
     assert (await security.read_tc_frame_counter(znp)) == 0x98765432
 
 
-def ieee_and_key(text):
+def ieee_and_key(text) -> zigpy.state.Key:
     ieee, key = text.replace(":", "").split("|")
 
-    return t.EUI64(bytes.fromhex(ieee)[::-1]), t.KeyData(bytes.fromhex(key))
+    return zigpy.state.Key(
+        partner_ieee=t.EUI64(bytes.fromhex(ieee)[::-1]),
+        key=t.KeyData(bytes.fromhex(key)),
+    )
 
 
 def test_seed_candidate_finding_simple():
-    ieee1, key1 = ieee_and_key("0011223344556677|000102030405060708090a0b0c0d0e0f")
-    ieee2, key2 = ieee_and_key("1111223344556677|101112131415161718191a1b1c1d1e1f")
+    k1 = ieee_and_key("0011223344556677|000102030405060708090a0b0c0d0e0f")
+    k2 = ieee_and_key("1111223344556677|101112131415161718191a1b1c1d1e1f")
 
-    (c1, s1), (c2, s2) = security.iter_seed_candidates([(ieee1, key1), (ieee2, key2)])
+    (count1, seed1), (count2, seed2) = security.iter_seed_candidates([k1, k2])
 
-    assert c1 == c2 == 1
+    assert count1 == count2 == 1
 
-    sh1 = security.find_key_shift(ieee1, key1, s1)
-    sh2 = security.find_key_shift(ieee2, key2, s2)
-    assert sh1 is not None and sh2 is not None
+    shift1 = security.find_key_shift(k1.partner_ieee, k1.key, seed1)
+    shift2 = security.find_key_shift(k2.partner_ieee, k2.key, seed2)
+    assert shift1 is not None and shift2 is not None
 
-    assert security.compute_key(ieee1, s1, sh1) == key1
-    assert security.compute_key(ieee2, s2, sh2) == key2
+    assert security.compute_key(k1.partner_ieee, seed1, shift1) == k1.key
+    assert security.compute_key(k2.partner_ieee, seed2, shift2) == k2.key
 
 
 def min_rotate(lst):
@@ -366,7 +369,7 @@ def test_seed_candidate_finding_complex():
     # One seed generated all but one of the keys, so there are 24 equally valid seeds.
     # They are all really rotations of the same seed.
     assert [c for c, s in candidates].count(24) == 24
-    assert len({min_rotate(s) for c, s in candidates if c == 24}) == 1
+    assert len({min_rotate(bytes(s)) for c, s in candidates if c == 24}) == 1
 
     # And one just for the bogus entry
     assert [c[0] for c in candidates].count(1) == 1

zigpy_znp/api.py

@@ -407,17 +407,25 @@ async def write_network_info(
 
         LOGGER.debug("Writing children and keys")
 
-        new_tclk_seed = await security.write_devices(
+        optimal_tclk_seed = security.find_optimal_tclk_seed(devices.values(), tclk_seed)
+
+        if tclk_seed != optimal_tclk_seed:
+            LOGGER.warning(
+                "Provided TCLK seed %s is not optimal, using %s instead.",
+                tclk_seed,
+                optimal_tclk_seed,
+            )
+
+            await self.nvram.osal_write(OsalNvIds.TCLK_SEED, optimal_tclk_seed)
+            tclk_seed = optimal_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/znp/security.py

@@ -28,17 +28,17 @@ def rotate(lst: typing.Sequence, n: int) -> typing.Sequence:
     return lst[n:] + lst[:n]
 
 
-def compute_key(ieee: t.EUI64, tclk_seed: bytes, shift: int) -> t.KeyData:
+def compute_key(ieee: t.EUI64, tclk_seed: t.KeyData, shift: int) -> t.KeyData:
     rotated_tclk_seed = rotate(tclk_seed, n=shift)
     return t.KeyData([a ^ b for a, b in zip(rotated_tclk_seed, 2 * ieee.serialize())])
 
 
-def compute_tclk_seed(ieee: t.EUI64, key: t.KeyData, shift: int) -> bytes:
+def compute_tclk_seed(ieee: t.EUI64, key: t.KeyData, shift: int) -> t.KeyData:
     rotated_tclk_seed = bytes(a ^ b for a, b in zip(key, 2 * ieee.serialize()))
-    return rotate(rotated_tclk_seed, n=-shift)
+    return t.KeyData(rotate(rotated_tclk_seed, n=-shift))
 
 
-def find_key_shift(ieee: t.EUI64, key: t.KeyData, tclk_seed: bytes) -> int | None:
+def find_key_shift(ieee: t.EUI64, key: t.KeyData, tclk_seed: t.KeyData) -> int | None:
     for shift in range(0x00, 0x0F + 1):
         if tclk_seed == compute_tclk_seed(ieee, key, shift):
             return shift
@@ -47,20 +47,26 @@ def find_key_shift(ieee: t.EUI64, key: t.KeyData, tclk_seed: bytes) -> int | Non
 
 
 def count_seed_matches(
-    ieees_and_keys: typing.Sequence[tuple[t.EUI64, t.KeyData]], tclk_seed: bytes
+    keys: typing.Sequence[zigpy.state.Key], tclk_seed: t.KeyData
 ) -> int:
-    return sum(find_key_shift(i, k, tclk_seed) is not None for i, k in ieees_and_keys)
+    count = 0
+
+    for key in keys:
+        if find_key_shift(key.partner_ieee, key.key, tclk_seed) is not None:
+            count += 1
+
+    return count
 
 
 def iter_seed_candidates(
-    ieees_and_keys: typing.Sequence[tuple[t.EUI64, t.KeyData]]
+    keys: typing.Sequence[zigpy.state.Key],
 ) -> typing.Iterable[tuple[int, t.KeyData]]:
-    for ieee, key in ieees_and_keys:
+    for key in keys:
         # Derive a seed from each candidate. All rotations of a seed are equivalent.
-        tclk_seed = t.KeyData(compute_tclk_seed(ieee, key, 0))
+        tclk_seed = compute_tclk_seed(key.partner_ieee, key.key, 0)
 
         # And see how many other keys share this same seed
-        count = count_seed_matches(ieees_and_keys, tclk_seed)
+        count = count_seed_matches(keys, tclk_seed)
 
         yield count, tclk_seed
 
@@ -167,7 +173,7 @@ async def read_addr_mgr_entries(znp: ZNP) -> typing.Sequence[t.AddrMgrEntry]:
 
 
 async def read_hashed_link_keys(
-    znp: ZNP, tclk_seed: bytes
+    znp: ZNP, tclk_seed: t.KeyData
 ) -> typing.Iterable[zigpy.state.Key]:
     if znp.version >= 3.30:
         entries = znp.nvram.read_table(
@@ -352,37 +358,31 @@ async def write_addr_manager_entries(
     await znp.nvram.osal_write(OsalNvIds.ADDRMGR, t.AddressManagerTable(new_entries))
 
 
+def find_optimal_tclk_seed(
+    devices: typing.Sequence[StoredDevice], tclk_seed: t.KeyData
+) -> t.KeyData:
+    keys = [d.key for d in devices if d.key]
+
+    if not keys:
+        return tclk_seed
+
+    best_count, best_seed = max(iter_seed_candidates(keys))
+    tclk_count = count_seed_matches(keys, tclk_seed)
+    assert tclk_count <= best_count
+
+    # Prefer the existing TCLK seed if it's as good as the others
+    if tclk_count == best_count:
+        return tclk_seed
+
+    return best_seed
+
+
 async def write_devices(
     znp: ZNP,
     devices: typing.Sequence[StoredDevice],
     counter_increment: t.uint32_t = 2500,
     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:
-        best_count, best_seed = max(iter_seed_candidates(ieees_and_keys))
-
-        # Check to see if the provided tclk_seed is also optimal
-        if tclk_seed is not None:
-            tclk_count = count_seed_matches(ieees_and_keys, tclk_seed)
-            assert tclk_count <= best_count
-
-            if tclk_count < best_count:
-                LOGGER.warning(
-                    "Provided TCLK seed %s only generates %d keys, but computed seed"
-                    " %s generates %d keys. Picking computed seed.",
-                    tclk_seed,
-                    tclk_count,
-                    best_seed,
-                    best_count,
-                )
-            else:
-                best_seed = tclk_seed
-
-        tclk_seed = best_seed
-
     hashed_link_key_table = []
     aps_key_data_table = []
     link_key_table = t.APSLinkKeyTable()
@@ -453,7 +453,7 @@ async def write_devices(
     await write_addr_manager_entries(znp, devices)
 
     if old_link_key_table is None:
-        return tclk_seed
+        return
 
     await znp.nvram.osal_write(OsalNvIds.APS_LINK_KEY_TABLE, new_link_key_table_value)
 
@@ -498,5 +498,3 @@ async def write_devices(
             values=aps_key_data_table,
             fill_value=aps_key_data_fill_value,
         )
-
-    return tclk_seed