8000 Return raw discovery result in cli discover raw (#1342) · python-kasa/python-kasa@bf8f0ad · GitHub
[go: up one dir, main page]
Skip to content

Commit bf8f0ad

Browse files
sdb9696sdb9696
authored
Return raw discovery result in cli discover raw (#1342)
Add `on_discovered_raw` callback to Discover and adds a cli command `discover raw` which returns the raw json before serializing to a `DiscoveryResult` and attempting to create a device class.
1 parent 464683e commit bf8f0ad

File tree

4 files changed

+158
-23
lines changed

4 files changed

+158
-23
lines changed

‎kasa/cli/discover.py

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,17 @@
1414
Discover,
1515
UnsupportedDeviceError,
1616
)
17-
from kasa.discover import ConnectAttempt, DiscoveryResult
17+
from kasa.discover import (
18+
NEW_DISCOVERY_REDACTORS,
19+
ConnectAttempt,
20+
DiscoveredRaw,
21+
DiscoveryResult,
22+
)
1823
from kasa.iot.iotdevice import _extract_sys_info
24+
from kasa.protocols.iotprotocol import REDACTORS as IOT_REDACTORS
25+
from kasa.protocols.protocol import redact_data
1926

27+
from ..json import dumps as json_dumps
2028
from .common import echo, error
2129

2230

@@ -64,7 +72,9 @@ async def print_discovered(dev: Device) -> None:
6472
await ctx.parent.invoke(state)
6573
echo()
6674

67-
discovered = await _discover(ctx, print_discovered, print_unsupported)
75+
discovered = await _discover(
76+
ctx, print_discovered=print_discovered, print_unsupported=print_unsupported
77+
)
6878
if ctx.parent.parent.params["host"]:
6979
return discovered
7080

@@ -77,6 +87,33 @@ async def print_discovered(dev: Device) -> None:
7787
return discovered
7888

7989

90+
@discover.command()
91+
@click.option(
92+
"--redact/--no-redact",
93+
default=False,
94+
is_flag=True,
95+
type=bool,
96+
help="Set flag to redact sensitive data from raw output.",
97+
)
98+
@click.pass_context
99+
async def raw(ctx, redact: bool):
100+
"""Return raw discovery data returned from devices."""
101+
102+
def print_raw(discovered: DiscoveredRaw):
103+
if redact:
104+
redactors = (
105+
NEW_DISCOVERY_REDACTORS
106+
if discovered["meta"]["port"] == Discover.DISCOVERY_PORT_2
107+
else IOT_REDACTORS
108+
)
109+
discovered["discovery_response"] = redact_data(
110+
discovered["discovery_response"], redactors
111+
)
112+
echo(json_dumps(discovered, indent=True))
113+
114+
return await _discover(ctx, print_raw=print_raw, do_echo=False)
115+
116+
80117
@discover.command()
81118
@click.pass_context
82119
async def list(ctx):
@@ -102,10 +139,17 @@ async def print_unsupported(unsupported_exception: UnsupportedDeviceError):
102139
echo(f"{host:<15} UNSUPPORTED DEVICE")
103140

104141
echo(f"{'HOST':<15} {'DEVICE FAMILY':<20} {'ENCRYPT':<7} {'ALIAS'}")
105-
return await _discover(ctx, print_discovered, print_unsupported, do_echo=False)
142+
return await _discover(
143+
ctx,
144+
print_discovered=print_discovered,
145+
print_unsupported=print_unsupported,
146+
do_echo=False,
147+
)
106148

107149

108-
async def _discover(ctx, print_discovered, print_unsupported, *, do_echo=True):
150+
async def _discover(
151+
ctx, *, print_discovered=None, print_unsupported=None, print_raw=None, do_echo=True
152+
):
109153
params = ctx.parent.parent.params
110154
target = params["target"]
111155
username = params["username"]
@@ -126,6 +170,7 @@ async def _discover(ctx, print_discovered, print_unsupported, *, do_echo=True):
126170
timeout=timeout,
127171
discovery_timeout=discovery_timeout,
128172
on_unsupported=print_unsupported,
173+
on_discovered_raw=print_raw,
129174
)
130175
if do_echo:
131176
echo(f"Discovering devices on {target} for {discovery_timeout} seconds")
@@ -137,6 +182,7 @@ async def _discover(ctx, print_discovered, print_unsupported, *, do_echo=True):
137182
port=port,
138183
timeout=timeout,
139184
credentials=credentials,
185+
on_discovered_raw=print_raw,
140186
)
141187

142188
for device in discovered_devices.values():

‎kasa/discover.py

Lines changed: 65 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@
9999
Annotated,
100100
Any,
101101
NamedTuple,
102+
TypedDict,
102103
cast,
103104
)
104105

@@ -147,18 +148,35 @@ class ConnectAttempt(NamedTuple):
147148
device: type
148149

149150

151+
class DiscoveredMeta(TypedDict):
152+
"""Meta info about discovery response."""
153+
154+
ip: str
155+
port: int
156+
157+
158+
class DiscoveredRaw(TypedDict):
159+
"""Try to connect attempt."""
160+
161+
meta: DiscoveredMeta
162+
discovery_response: dict
163+
164+
150165
OnDiscoveredCallable = Callable[[Device], Coroutine]
166+
OnDiscoveredRawCallable = Callable[[DiscoveredRaw], None]
151167
OnUnsupportedCallable = Callable[[UnsupportedDeviceError], Coroutine]
152168
OnConnectAttemptCallable = Callable[[ConnectAttempt, bool], None]
153169
DeviceDict = dict[str, Device]
154170

155171
NEW_DISCOVERY_REDACTORS: dict[str, Callable[[Any], Any] | None] = {
156172
"device_id": lambda x: "REDACTED_" + x[9::],
173+
"device_name": lambda x: "#MASKED_NAME#" if x else "",
157174
"owner": lambda x: "REDACTED_" + x[9::],
158175
"mac": mask_mac,
159176
"master_device_id": lambda x: "REDACTED_" + x[9::],
160177
"group_id": lambda x: "REDACTED_" + x[9::],
161178
"group_name": lambda x: "I01BU0tFRF9TU0lEIw==",
179+
"encrypt_info": lambda x: {**x, "key": "", "data": ""},
162180
}
163181

164182

@@ -216,6 +234,7 @@ def __init__(
216234
self,
217235
*,
218236
on_discovered: OnDiscoveredCallable | None = None,
237+
on_discovered_raw: OnDiscoveredRawCallable | None = None,
219238
target: str = "255.255.255.255",
220239
discovery_packets: int = 3,
221240
discovery_timeout: int = 5,
@@ -240,6 +259,7 @@ def __init__(
240259
self.unsupported_device_exceptions: dict = {}
241260
self.invalid_device_exceptions: dict = {}
242261
self.on_unsupported = on_unsupported
262+
self.on_discovered_raw = on_discovered_raw
243263
self.credentials = credentials
244264
self.timeout = timeout
245265
self.discovery_timeout = discovery_timeout
@@ -329,12 +349,23 @@ def datagram_received(
329349
config.timeout = self.timeout
330350
try:
331351
if port == self.discovery_port:
332-
device = Discover._get_device_instance_legacy(data, config)
352+
json_func = Discover._get_discovery_json_legacy
353+
device_func = Discover._get_device_instance_legacy
333354
elif port == Discover.DISCOVERY_PORT_2:
334355
config.uses_http = True
335-
device = Discover._get_device_instance(data, config)
356+
json_func = Discover._get_discovery_json
357+
device_func = Discover._get_device_instance
336358
else:
337359
return
360+
info = json_func(data, ip)
361+
if self.on_discovered_raw is not None:
362+
self.on_discovered_raw(
363+
{
364+
"discovery_response": info,
365+
"meta": {"ip": ip, "port": port},
366+
}
367+
)
368+
device = device_func(info, config)
338369
except UnsupportedDeviceError as udex:
339370
_LOGGER.debug("Unsupported device found at %s << %s", ip, udex)
340371
self.unsupported_device_exceptions[ip] = udex
@@ -391,6 +422,7 @@ async def discover(
391422
*,
392423
target: str = "255.255.255.255",
393424
on_discovered: OnDiscoveredCallable | None = None,
425+
on_discovered_raw: OnDiscoveredRawCallable | None = None,
394426
discovery_timeout: int = 5,
395427
discovery_packets: int = 3,
396428
interface: str | None = None,
@@ -421,6 +453,8 @@ async def discover(
421453
:param target: The target address where to send the broadcast discovery
422454
queries if multi-homing (e.g. 192.168.xxx.255).
423455
:param on_discovered: coroutine to execute on discovery
456+
:param on_discovered_raw: Optional callback once discovered json is loaded
457+
before any attempt to deserialize it and create devices
424458
:param discovery_timeout: Seconds to wait for responses, defaults to 5
425459
:param discovery_packets: Number of discovery packets to broadcast
426460
:param interface: Bind to specific interface
@@ -443,6 +477,7 @@ async def discover(
443477
discovery_packets=discovery_packets,
444478
interface=interface,
445479
on_unsupported=on_unsupported,
480+
on_discovered_raw=on_discovered_raw,
446481
credentials=credentials,
447482
timeout=timeout,
448483
discovery_timeout=discovery_timeout,
@@ -476,6 +511,7 @@ async def discover_single(
476511
credentials: Credentials | None = None,
477512
username: str | None = None,
478513
password: str | None = None,
514+
on_discovered_raw: OnDiscoveredRawCallable | None = None,
479515
on_unsupported: OnUnsupportedCallable | None = None,
480516
) -> Device | None:
481517
"""Discover a single device by the given IP address.
@@ -493,6 +529,9 @@ async def discover_single(
493529
username and password are ignored if provided.
494530
:param username: Username for devices that require authentication
495531
:param password: Password for devices that require authentication
532+
:param on_discovered_raw: Optional callback once discovered json is loaded
533+
before any attempt to deserialize it and create devices
534+
:param on_unsupported: Optional callback when unsupported devices are discovered
496535
:rtype: SmartDevice
497536
:return: Object for querying/controlling found device.
498537
"""
@@ -529,6 +568,7 @@ async def discover_single(
529568
credentials=credentials,
530569
timeout=timeout,
531570
discovery_timeout=discovery_timeout,
571+
on_discovered_raw=on_discovered_raw,
532572
),
533573
local_addr=("0.0.0.0", 0), # noqa: S104
534574
)
@@ -666,15 +706,19 @@ def _get_device_class(info: dict) -> type[Device]:
666706
return get_device_class_from_sys_info(info)
667707

668708
@staticmethod
669-
def _get_device_instance_legacy(data: bytes, config: DeviceConfig) -> IotDevice:
670-
"""Get SmartDevice from legacy 9999 response."""
709+
def _get_discovery_json_legacy(data: bytes, ip: str) -> dict:
710+
"""Get discovery json from legacy 9999 response."""
671711
try:
672712
info = json_loads(XorEncryption.decrypt(data))
673713
except Exception as ex:
674714
raise KasaException(
675-
f"Unable to read response from device: {config.host}: {ex}"
715+
f"Unable to read response from device: {ip}: {ex}"
676716
) from ex
717+
return info
677718

719+
@staticmethod
720+
def _get_device_instance_legacy(info: dict, config: DeviceConfig) -> Device:
721+
"""Get IotDevice from legacy 9999 response."""
678722
if _LOGGER.isEnabledFor(logging.DEBUG):
679723
data = redact_data(info, IOT_REDACTORS) if Discover._redact_data else info
680724
_LOGGER.debug("[DISCOVERY] %s << %s", config.host, pf(data))
@@ -716,19 +760,24 @@ def _decrypt_discovery_data(discovery_result: DiscoveryResult) -> None:
716760
discovery_result.decrypted_data = json_loads(decrypted_data)
717761

718762
@staticmethod
719-
def _get_device_instance(
720-
data: bytes,
721-
config: DeviceConfig,
722-
) -> Device:
723-
"""Get SmartDevice from the new 20002 response."""
724-
debug_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
763+
def _get_discovery_json(data: bytes, ip: str) -> dict:
764+
"""Get discovery json from the new 20002 response."""
725765
try:
726766
info = json_loads(data[16:])
727767
except Exception as ex:
728-
_LOGGER.debug("Got invalid response from device %s: %s", config.host, data)
768+
_LOGGER.debug("Got invalid response from device %s: %s", ip, data)
729769
raise KasaException(
730-
f"Unable to read response from device: {config.host}: {ex}"
770+
f"Unable to read response from device: {ip}: {ex}"
731771
) from ex
772+
return info
773+
774+
@staticmethod
775+
def _get_device_instance(
776+
info: dict,
777+
config: DeviceConfig,
778+
) -> Device:
779+
"""Get SmartDevice from the new 20002 response."""
780+
debug_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
732781

733782
try:
734783
discovery_result = DiscoveryResult.from_dict(info["result"])
@@ -757,7 +806,9 @@ def _get_device_instance(
757806
Discover._decrypt_discovery_data(discovery_result)
758807
except Exception:
759808
_LOGGER.exception(
760-
"Unable to decrypt discovery data %s: %s", config.host, data
809+
"Unable to decrypt discovery data %s: %s",
810+
config.host,
811+
redact_data(info, NEW_DISCOVERY_REDACTORS),
761812
)
762813

763814
type_ = discovery_result.device_type

‎kasa/json.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,18 +8,24 @@
88
try:
99
import orjson
1010

11-
def dumps(obj: Any, *, default: Callable | None = None) -> str:
11+
def dumps(
12+
obj: Any, *, default: Callable | None = None, indent: bool = False
13+
) -> str:
1214
"""Dump JSON."""
13-
return orjson.dumps(obj).decode()
15+
return orjson.dumps(
16+
obj, option=orjson.OPT_INDENT_2 if indent else None
17+
).decode()
1418

1519
loads = orjson.loads
1620
except ImportError:
1721
import json
1822

19-
def dumps(obj: Any, *, default: Callable | None = None) -> str:
23+
def dumps(
24+
obj: Any, *, default: Callable | None = None, indent: bool = False
25+
) -> str:
2026
"""Dump JSON."""
2127
# Separators specified for consistency with orjson
22-
return json.dumps(obj, separators=(",", ":"))
28+
return json.dumps(obj, separators=(",", ":"), indent=2 if indent else None)
2329

2430
loads = json.loads
2531

‎tests/test_cli.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,9 @@
4242
from kasa.cli.time import time
4343
from kasa.cli.usage import energy
4444
from kasa.cli.wifi import wifi
45-
from kasa.discover import Discover, DiscoveryResult
45+
from kasa.discover import Discover, DiscoveryResult, redact_data
4646
from kasa.iot import IotDevice
47+
from kasa.json import dumps as json_dumps
4748
from kasa.smart import SmartDevice
4849
from kasa.smartcam import SmartCamDevice
4950

@@ -126,6 +127,36 @@ async def test_list_devices(discovery_mock, runner):
126127
assert row in res.output
127128

128129

130+
async def test_discover_raw(discovery_mock, runner, mocker):
131+
"""Test the discover raw command."""
132+
redact_spy = mocker.patch(
133+
"kasa.protocols.protocol.redact_data", side_effect=redact_data
134+
)
135+
res = await runner.invoke(
136+
cli,
137+
["--username", "foo", "--password", "bar", "discover", "raw"],
138+
catch_exceptions=False,
139+
)
140+
assert res.exit_code == 0
141+
142+
expected = {
143+
"discovery_response": discovery_mock.discovery_data,
144+
"meta": {"ip": "127.0.0.123", "port": discovery_mock.discovery_port},
145+
}
146+
assert res.output == json_dumps(expected, indent=True) + "\n"
147+
148+
redact_spy.assert_not_called()
149+
150+
res = await runner.invoke(
151+
cli,
152+
["--username", "foo", "--password", "bar", "discover", "raw", "--redact"],
153+
catch_exceptions=False,
154+
)
155+
assert res.exit_code == 0
156+
157+
redact_spy.assert_called()
158+
159+
129160
@new_discovery
130161
async def test_list_auth_failed(discovery_mock, mocker, runner):
131162
"""Test that device update is called on main."""
@@ -731,6 +762,7 @@ async def test_without_device_type(dev, mocker, runner):
731762
timeout=5,
732763
discovery_timeout=7,
733764
on_unsupported=ANY,
765+
on_discovered_raw=ANY,
734766
)
735767

736768

0 commit comments

Comments
 (0)
0