Update dump_devinfo for raw discovery json and common redactors (#1358)

This PR does a few related things to dump_devinfo:

- Store the raw discovery result in the fixture.
- Consolidate redaction logic so it's not duplicated in dump_devinfo.
- Update existing fixtures to:
  - Store raw discovery result under `result`
  - Use `SCRUBBED_CHILD_DEVICE_ID` everywhere
  - Have correct values as per the consolidated redactors.

devtools/dump_devinfo.py

@@ -10,8 +10,6 @@ and finally execute a query to query all of them at once.
 
 from __future__ import annotations
 
-import base64
-import collections.abc
 import dataclasses
 import json
 import logging
@@ -19,6 +17,7 @@ import re
 import sys
 import traceback
 from collections import defaultdict, namedtuple
+from collections.abc import Callable
 from pathlib import Path
 from pprint import pprint
 from typing import Any
@@ -39,13 +38,20 @@ from kasa import (
 )
 from kasa.device_factory import get_protocol
 from kasa.deviceconfig import DeviceEncryptionType, DeviceFamily
-from kasa.discover import DiscoveryResult
+from kasa.discover import (
+    NEW_DISCOVERY_REDACTORS,
+    DiscoveredRaw,
+    DiscoveryResult,
+)
 from kasa.exceptions import SmartErrorCode
 from kasa.protocols import IotProtocol
+from kasa.protocols.iotprotocol import REDACTORS as IOT_REDACTORS
+from kasa.protocols.protocol import redact_data
 from kasa.protocols.smartcamprotocol import (
     SmartCamProtocol,
     _ChildCameraProtocolWrapper,
 )
+from kasa.protocols.smartprotocol import REDACTORS as SMART_REDACTORS
 from kasa.protocols.smartprotocol import SmartProtocol, _ChildProtocolWrapper
 from kasa.smart import SmartChildDevice, SmartDevice
 from kasa.smartcam import SmartCamDevice
@@ -63,6 +69,42 @@ ENCRYPT_TYPES = [encrypt_type.value for encrypt_type in DeviceEncryptionType]
 _LOGGER = logging.getLogger(__name__)
 
 
+def _wrap_redactors(redactors: dict[str, Callable[[Any], Any] | None]):
+    """Wrap the redactors for  dump_devinfo.
+
+    Will replace all partial REDACT_ values with zeros.
+    If the data item is already scrubbed by dump_devinfo will leave as-is.
+    """
+
+    def _wrap(key: str) -> Any:
+        def _wrapped(redactor: Callable[[Any], Any] | None) -> Any | None:
+            if redactor is None:
+                return lambda x: "**SCRUBBED**"
+
+            def _redact_to_zeros(x: Any) -> Any:
+                if isinstance(x, str) and "REDACT" in x:
+                    return re.sub(r"\w", "0", x)
+                if isinstance(x, dict):
+                    for k, v in x.items():
+                        x[k] = _redact_to_zeros(v)
+                return x
+
+            def _scrub(x: Any) -> Any:
+                if key in {"ip", "local_ip"}:
+                    return "127.0.0.123"
+                # Already scrubbed by dump_devinfo
+                if isinstance(x, str) and "SCRUBBED" in x:
+                    return x
+                default = redactor(x)
+                return _redact_to_zeros(default)
+
+            return _scrub
+
+        return _wrapped(redactors[key])
+
+    return {key: _wrap(key) for key in redactors}
+
+
 @dataclasses.dataclass
 class SmartCall:
     """Class for smart and smartcam calls."""
@@ -74,115 +116,6 @@ class SmartCall:
     supports_multiple: bool = True
 
 
-def scrub(res):
-    """Remove identifiers from the given dict."""
-    keys_to_scrub = [
-        "deviceId",
-        "fwId",
-        "hwId",
-        "oemId",
-        "mac",
-        "mic_mac",
-        "latitude_i",
-        "longitude_i",
-        "latitude",
-        "longitude",
-        "la",  # lat on ks240
-        "lo",  # lon on ks240
-        "owner",
-        "device_id",
-        "ip",
-        "ssid",
-        "hw_id",
-        "fw_id",
-        "oem_id",
-        "nickname",
-        "alias",
-        "bssid",
-        "channel",
-        "original_device_id",  # for child devices on strips
-        "parent_device_id",  # for hub children
-        "setup_code",  # matter
-        "setup_payload",  # matter
-        "mfi_setup_code",  # mfi_ for homekit
-        "mfi_setup_id",
-        "mfi_token_token",
-        "mfi_token_uuid",
-        "dev_id",
-        "device_name",
-        "device_alias",
-        "connect_ssid",
-        "encrypt_info",
-        "local_ip",
-        "username",
-        # vacuum
-        "board_sn",
-        "custom_sn",
-        "location",
-    ]
-
-    for k, v in res.items():
-        if isinstance(v, collections.abc.Mapping):
-            if k == "encrypt_info":
-                if "data" in v:
-                    v["data"] = ""
-                if "key" in v:
-                    v["key"] = ""
-            else:
-                res[k] = scrub(res.get(k))
-        elif (
-            isinstance(v, list)
-            and len(v) > 0
-            and isinstance(v[0], collections.abc.Mapping)
-        ):
-            res[k] = [scrub(vi) for vi in v]
-        else:
-            if k in keys_to_scrub:
-                if k in ["mac", "mic_mac"]:
-                    # Some macs have : or - as a separator and others do not
-                    if len(v) == 12:
-                        v = f"{v[:6]}000000"
-                    else:
-                        delim = ":" if ":" in v else "-"
-                        rest = delim.join(
-                            format(s, "02x") for s in bytes.fromhex("000000")
-                        )
-                        v = f"{v[:8]}{delim}{rest}"
-                elif k in ["latitude", "latitude_i", "longitude", "longitude_i"]:
-                    v = 0
-                elif k in ["ip", "local_ip"]:
-                    v = "127.0.0.123"
-                elif k in ["ssid"]:
-                    # Need a valid base64 value here
-                    v = base64.b64encode(b"#MASKED_SSID#").decode()
-                elif k in ["nickname"]:
-                    v = base64.b64encode(b"#MASKED_NAME#").decode()
-                elif k in [
-                    "alias",
-                    "device_alias",
-                    "device_name",
-                    "username",
-                    "location",
-                ]:
-                    v = "#MASKED_NAME#"
-                elif isinstance(res[k], int):
-                    v = 0
-                elif k in ["map_data"]:  #
-                    v = "#SCRUBBED_MAPDATA#"
-                elif k in ["device_id", "dev_id"] and "SCRUBBED" in v:
-                    pass  # already scrubbed
-                elif k == ["device_id", "dev_id"] and len(v) > 40:
-                    # retain the last two chars when scrubbing child ids
-                    end = v[-2:]
-                    v = re.sub(r"\w", "0", v)
-                    v = v[:40] + end
-                else:
-                    v = re.sub(r"\w", "0", v)
-
-            res[k] = v
-    return res
-
-
 def default_to_regular(d):
     """Convert nested defaultdicts to regular ones.
 
@@ -209,7 +142,7 @@ async def handle_device(
     for fixture_result in fixture_results:
         save_filename = Path(basedir) / fixture_result.folder / fixture_result.filename
 
-        pprint(scrub(fixture_result.data))
+        pprint(fixture_result.data)
         if autosave:
             save = "y"
         else:
@@ -325,6 +258,11 @@ async def cli(
     if debug:
         logging.basicConfig(level=logging.DEBUG)
 
+    raw_discovery = {}
+
+    def capture_raw(discovered: DiscoveredRaw):
+        raw_discovery[discovered["meta"]["ip"]] = discovered["discovery_response"]
+
     credentials = Credentials(username=username, password=password)
     if host is not None:
         if discovery_info:
@@ -377,12 +315,16 @@ async def cli(
                 credentials=credentials,
                 port=port,
                 discovery_timeout=discovery_timeout,
+                on_discovered_raw=capture_raw,
             )
+            discovery_info = raw_discovery[device.host]
+            if decrypted_data := device._discovery_info.get("decrypted_data"):
+                discovery_info["decrypted_data"] = decrypted_data
             await handle_device(
                 basedir,
                 autosave,
                 device.protocol,
-                discovery_info=device._discovery_info,
+                discovery_info=discovery_info,
                 batch_size=batch_size,
             )
     else:
@@ -391,21 +333,28 @@ async def cli(
             f" {target}. Use --target to override."
         )
         devices = await Discover.discover(
-            target=target, credentials=credentials, discovery_timeout=discovery_timeout
+            target=target,
+            credentials=credentials,
+            discovery_timeout=discovery_timeout,
+            on_discovered_raw=capture_raw,
         )
         click.echo(f"Detected {len(devices)} devices")
         for dev in devices.values():
+            discovery_info = raw_discovery[dev.host]
+            if decrypted_data := dev._discovery_info.get("decrypted_data"):
+                discovery_info["decrypted_data"] = decrypted_data
+
             await handle_device(
                 basedir,
                 autosave,
                 dev.protocol,
-                discovery_info=dev._discovery_info,
+                discovery_info=discovery_info,
                 batch_size=batch_size,
             )
 
 
 async def get_legacy_fixture(
-    protocol: IotProtocol, *, discovery_info: dict[str, Any] | None
+    protocol: IotProtocol, *, discovery_info: dict[str, dict[str, Any]] | None
 ) -> FixtureResult:
     """Get fixture for legacy IOT style protocol."""
     items = [
@@ -475,11 +424,21 @@ async def get_legacy_fixture(
         _echo_error(f"Unable to query all successes at once: {ex}")
     finally:
         await protocol.close()
+
+    final = redact_data(final, _wrap_redactors(IOT_REDACTORS))
+
+    # Scrub the child device ids
+    if children := final.get("system", {}).get("get_sysinfo", {}).get("children"):
+        for index, child in enumerate(children):
+            if "id" not in child:
+                _LOGGER.error("Could not find a device for the child device: %s", child)
+            else:
+                child["id"] = f"SCRUBBED_CHILD_DEVICE_ID_{index + 1}"
+
     if discovery_info and not discovery_info.get("system"):
-        # Need to recreate a DiscoverResult here because we don't want the aliases
-        # in the fixture, we want the actual field names as returned by the device.
-        dr = DiscoveryResult.from_dict(discovery_info)
-        final["discovery_result"] = dr.to_dict()
+        final["discovery_result"] = redact_data(
+            discovery_info, _wrap_redactors(NEW_DISCOVERY_REDACTORS)
+        )
 
     click.echo(f"Got {len(successes)} successes")
     click.echo(click.style("## device info file ##", bold=True))
@@ -867,7 +826,10 @@ def get_smart_child_fixture(response):
 
 
 async def get_smart_fixtures(
-    protocol: SmartProtocol, *, discovery_info: dict[str, Any] | None, batch_size: int
+    protocol: SmartProtocol,
+    *,
+    discovery_info: dict[str, dict[str, Any]] | None,
+    batch_size: int,
 ) -> list[FixtureResult]:
     """Get fixture for new TAPO style protocol."""
     if isinstance(protocol, SmartCamProtocol):
@@ -988,22 +950,24 @@ async def get_smart_fixtures(
                 continue
             _LOGGER.error("Could not find a device for the child device: %s", child)
 
-    # Need to recreate a DiscoverResult here because we don't want the aliases
-    # in the fixture, we want the actual field names as returned by the device.
+    final = redact_data(final, _wrap_redactors(SMART_REDACTORS))
+    discovery_result = None
     if discovery_info:
-        dr = DiscoveryResult.from_dict(discovery_info)  # type: ignore
-        final["discovery_result"] = dr.to_dict()
+        final["discovery_result"] = redact_data(
+            discovery_info, _wrap_redactors(NEW_DISCOVERY_REDACTORS)
+        )
+        discovery_result = discovery_info["result"]
 
     click.echo(f"Got {len(successes)} successes")
     click.echo(click.style("## device info file ##", bold=True))
 
     if "get_device_info" in final:
         # smart protocol
-        model_info = SmartDevice._get_device_info(final, discovery_info)
+        model_info = SmartDevice._get_device_info(final, discovery_result)
         copy_folder = SMART_FOLDER
     else:
         # smart camera protocol
-        model_info = SmartCamDevice._get_device_info(final, discovery_info)
+        model_info = SmartCamDevice._get_device_info(final, discovery_result)
         copy_folder = SMARTCAM_FOLDER
     hw_version = model_info.hardware_version
     sw_version = model_info.firmware_version

devtools/generate_supported.py

@@ -205,7 +205,7 @@ def _get_supported_devices(
             fixture_data = json.load(f)
 
         model_info = device_cls._get_device_info(
-            fixture_data, fixture_data.get("discovery_result")
+            fixture_data, fixture_data.get("discovery_result", {}).get("result")
         )
 
         supported_type = DEVICE_TYPE_TO_PRODUCT_GROUP[model_info.device_type]

devtools/update_fixtures.py

@@ -0,0 +1,128 @@
+"""Module to mass update fixture files."""
+
+import json
+import logging
+from collections.abc import Callable
+from pathlib import Path
+
+import asyncclick as click
+
+from devtools.dump_devinfo import _wrap_redactors
+from kasa.discover import NEW_DISCOVERY_REDACTORS, redact_data
+from kasa.protocols.iotprotocol import REDACTORS as IOT_REDACTORS
+from kasa.protocols.smartprotocol import REDACTORS as SMART_REDACTORS
+
+FIXTURE_FOLDER = "tests/fixtures/"
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def update_fixtures(update_func: Callable[[dict], bool], *, dry_run: bool) -> None:
+    """Run the update function against the fixtures."""
+    for file in Path(FIXTURE_FOLDER).glob("**/*.json"):
+        with file.open("r") as f:
+            fixture_data = json.load(f)
+
+        if file.parent.name == "serialization":
+            continue
+        changed = update_func(fixture_data)
+        if changed:
+            click.echo(f"Will update {file.name}\n")
+        if changed and not dry_run:
+            with file.open("w") as f:
+                json.dump(fixture_data, f, sort_keys=True, indent=4)
+                f.write("\n")
+
+
+def _discovery_result_update(info) -> bool:
+    """Update discovery_result to be the raw result and error_code."""
+    if (disco_result := info.get("discovery_result")) and "result" not in disco_result:
+        info["discovery_result"] = {
+            "result": disco_result,
+            "error_code": 0,
+        }
+        return True
+    return False
+
+
+def _child_device_id_update(info) -> bool:
+    """Update child device ids to be the scrubbed ids from dump_devinfo."""
+    changed = False
+    if get_child_device_list := info.get("get_child_device_list"):
+        child_device_list = get_child_device_list["child_device_list"]
+        child_component_list = info["get_child_device_component_list"][
+            "child_component_list"
+        ]
+        for index, child_device in enumerate(child_device_list):
+            child_component = child_component_list[index]
+            if "SCRUBBED" not in child_device["device_id"]:
+                dev_id = f"SCRUBBED_CHILD_DEVICE_ID_{index + 1}"
+                click.echo(
+                    f"child_device_id{index}: {child_device['device_id']} -> {dev_id}"
+                )
+                child_device["device_id"] = dev_id
+                child_component["device_id"] = dev_id
+                changed = True
+
+    if children := info.get("system", {}).get("get_sysinfo", {}).get("children"):
+        for index, child_device in enumerate(children):
+            if "SCRUBBED" not in child_device["id"]:
+                dev_id = f"SCRUBBED_CHILD_DEVICE_ID_{index + 1}"
+                click.echo(f"child_device_id{index}: {child_device['id']} -> {dev_id}")
+                child_device["id"] = dev_id
+                changed = True
+
+    return changed
+
+
+def _diff_data(fullkey, data1, data2, diffs):
+    if isinstance(data1, dict):
+        for k, v in data1.items():
+            _diff_data(fullkey + "/" + k, v, data2[k], diffs)
+    elif isinstance(data1, list):
+        for index, item in enumerate(data1):
+            _diff_data(fullkey + "/" + str(index), item, data2[index], diffs)
+    elif data1 != data2:
+        diffs[fullkey] = (data1, data2)
+
+
+def _redactor_result_update(info) -> bool:
+    """Update fixtures with the output using the common redactors."""
+    changed = False
+
+    redactors = IOT_REDACTORS if "system" in info else SMART_REDACTORS
+
+    for key, val in info.items():
+        if not isinstance(val, dict):
+            continue
+        if key == "discovery_result":
+            info[key] = redact_data(val, _wrap_redactors(NEW_DISCOVERY_REDACTORS))
+        else:
+            info[key] = redact_data(val, _wrap_redactors(redactors))
+        diffs: dict[str, tuple[str, str]] = {}
+        _diff_data(key, val, info[key], diffs)
+        if diffs:
+            for k, v in diffs.items():
+                click.echo(f"{k}: {v[0]} -> {v[1]}")
+            changed = True
+
+    return changed
+
+
+@click.option(
+    "--dry-run/--no-dry-run",
+    default=False,
+    is_flag=True,
+    type=bool,
+    help="Perform a dry run without saving.",
+)
+@click.command()
+async def cli(dry_run: bool) -> None:
+    """Cli method fo rupdating fixtures."""
+    update_fixtures(_discovery_result_update, dry_run=dry_run)
+    update_fixtures(_child_device_id_update, dry_run=dry_run)
+    update_fixtures(_redactor_result_update, dry_run=dry_run)
+
+
+if __name__ == "__main__":
+    cli()