FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Convert discovery to asyncio by rytilahti · Pull Request #13 · python-kasa/python-kasa · GitHub

Repository navigation

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .py  (2) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
178 changes: 126 additions & 52 deletions kasa/discover.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
"""Discovery module for TP-Link Smart Home devices."""
import asyncio
import json
import logging
import socket
from typing import Dict, Optional, Type
from typing import Dict, Mapping, Optional, Type, Union, cast

from kasa.protocol import TPLinkSmartHomeProtocol
from kasa.smartbulb import SmartBulb
Expand All @@ -13,6 +14,74 @@
_LOGGER = logging.getLogger(__name__)


class _DiscoverProtocol(asyncio.DatagramProtocol):
"""Implementation of the discovery protocol handler.

This is internal class, use :func:Discover.discover: instead.
"""

discovered_devices: Dict[str, SmartDevice]
discovered_devices_raw: Dict[str, Dict]

def __init__(
self,
*,
on_discovered=None,
target="255.255.255.255",
timeout=5,
discovery_packets=3,
):
self.transport = None
self.tries = discovery_packets
self.timeout = timeout
self.on_discovered = on_discovered
self.protocol = TPLinkSmartHomeProtocol()
self.target = (target, Discover.DISCOVERY_PORT)
self.discovered_devices = {}
self.discovered_devices_raw = {}

def connection_made(self, transport):
self.transport = transport
sock = transport.get_extra_info("socket")
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

self.do_discover()

def do_discover(self):
req = json.dumps(Discover.DISCOVERY_QUERY)
_LOGGER.debug("[DISCOVERY] %s >> %s", self.target, Discover.DISCOVERY_QUERY)
encrypted_req = self.protocol.encrypt(req)
for i in range(self.tries):
self.transport.sendto(encrypted_req[4:], self.target)

def datagram_received(self, data, addr):
ip, port = addr
if ip in self.discovered_devices:
return

info = json.loads(self.protocol.decrypt(data))
_LOGGER.debug("[DISCOVERY] %s << %s", ip, info)

device_class = Discover._get_device_class(info)
device = device_class(ip)

self.discovered_devices[ip] = device
self.discovered_devices_raw[ip] = info

if device_class is not None:
if self.on_discovered is not None:
asyncio.ensure_future(self.on_discovered(device))
else:
_LOGGER.error("Received invalid response: %s", info)

def error_received(self, ex):
_LOGGER.error("Got error: %s", ex)

def connection_lost(self, ex):
pass


class Discover:
"""Discover TPLink Smart Home devices.

Expand All @@ -28,6 +97,8 @@ class Discover:
The protocol uses UDP broadcast datagrams on port 9999 for discovery.
"""

DISCOVERY_PORT = 9999

DISCOVERY_QUERY = {
"system": {"get_sysinfo": None},
"emeter": {"get_realtime": None},
Expand All @@ -37,86 +108,76 @@ class Discover:
}

@staticmethod
def discover(
protocol: TPLinkSmartHomeProtocol = None,
target: str = "255.255.255.255",
port: int = 9999,
timeout: int = 3,
async def discover(
*,
target="255.255.255.255",
on_discovered=None,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

A type hint might be useful here.

OTOH it's a private class.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

Indeed; this wasn't completely done. I'll make the changes in the other PR considering it already contains all of this. So I'm closing this now, thanks for the review!

timeout=5,
discovery_packets=3,
return_raw=False,
) -> Dict[str, SmartDevice]:
"""Discover devices.
) -> Mapping[str, Union[SmartDevice, Dict]]:
"""Discover supported devices.

Sends discovery message to 255.255.255.255:9999 in order
to detect available supported devices in the local network,
and waits for given timeout for answers from devices.

:param protocol: Protocol implementation to use
If given, `on_discovered` coroutine will get passed with the SmartDevice as parameter.
The results of the discovery can be accessed either via `discovered_devices` (SmartDevice-derived) or
`discovered_devices_raw` (JSON objects).

:param target: The target broadcast address (e.g. 192.168.xxx.255).
:param timeout: How long to wait for responses, defaults to 3
:param port: port to send broadcast messages, defaults to 9999.
:rtype: dict
:return: Array of json objects {"ip", "port", "sys_info"}
:param on_discovered:
:param timeout: How long to wait for responses, defaults to 5
:param discovery_packets: Number of discovery packets are broadcasted.
:param return_raw: True to return JSON objects instead of Devices.
:return:
"""
if protocol is None:
protocol = TPLinkSmartHomeProtocol()
loop = asyncio.get_event_loop()
transport, protocol = await loop.create_datagram_endpoint(
lambda: _DiscoverProtocol(
target=target,
on_discovered=on_discovered,
timeout=timeout,
discovery_packets=discovery_packets,
),
local_addr=("0.0.0.0", 0),
)
protocol = cast(_DiscoverProtocol, protocol)

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.settimeout(timeout)

req = json.dumps(Discover.DISCOVERY_QUERY)
_LOGGER.debug("Sending discovery to %s:%s", target, port)
try:
_LOGGER.debug("Waiting %s seconds for responses...", timeout)
await asyncio.sleep(5)
finally:
transport.close()

encrypted_req = protocol.encrypt(req)
for i in range(discovery_packets):
sock.sendto(encrypted_req[4:], (target, port))
_LOGGER.debug("Discovered %s devices", len(protocol.discovered_devices))

devices = {}
_LOGGER.debug("Waiting %s seconds for responses...", timeout)
if return_raw:
return protocol.discovered_devices_raw

try:
while True:
data, addr = sock.recvfrom(4096)
ip, port = addr
info = json.loads(protocol.decrypt(data))
device_class = Discover._get_device_class(info)
if return_raw:
devices[ip] = info
elif device_class is not None:
devices[ip] = device_class(ip)
except socket.timeout:
_LOGGER.debug("Got socket timeout, which is okay.")
except Exception as ex:
_LOGGER.error("Got exception %s", ex, exc_info=True)
_LOGGER.debug("Found %s devices: %s", len(devices), devices)
return devices
return protocol.discovered_devices

@staticmethod
async def discover_single(
host: str, protocol: TPLinkSmartHomeProtocol = None
) -> Optional[SmartDevice]:
async def discover_single(host: str) -> Optional[SmartDevice]:
"""Discover a single device by the given IP address.

:param host: Hostname of device to query
:param protocol: Protocol implementation to use
:rtype: SmartDevice
:return: Object for querying/controlling found device.
"""
if protocol is None:
protocol = TPLinkSmartHomeProtocol()
protocol = TPLinkSmartHomeProtocol()

info = await protocol.query(host, Discover.DISCOVERY_QUERY)

device_class = Discover._get_device_class(info)
if device_class is not None:
return device_class(host)

return None
raise SmartDeviceException("Unable to discover device, received: %s" % info)

@staticmethod
def _get_device_class(info: dict) -> Optional[Type[SmartDevice]]:
def _get_device_class(info: dict) -> Type[SmartDevice]:
"""Find SmartDevice subclass for device described by passed data."""
if "system" in info and "get_sysinfo" in info["system"]:
sysinfo = info["system"]["get_sysinfo"]
Expand All @@ -136,4 +197,17 @@ def _get_device_class(info: dict) -> Optional[Type[SmartDevice]]:
elif "smartbulb" in type_.lower():
return SmartBulb

return None
raise SmartDeviceException("Unknown device type: %s", type_)


if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
loop = asyncio.get_event_loop()

async def _on_device(dev):
await dev.update()
_LOGGER.info("Got device: %s", dev)

devices = loop.run_until_complete(Discover.discover(on_discovered=_on_device))
for ip, dev in devices.items():
print(f"[{ip}] {dev}")
3 changes: 1 addition & 2 deletions kasa/smartdevice.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -646,14 +646,13 @@ def is_variable_color_temp(self) -> bool:
return False

def __repr__(self):
self.sync.update()
return "<{} model {} at {} ({}), is_on: {} - dev specific: {}>".format(
self.__class__.__name__,
self.model,
self.host,
self.alias,
self.is_on,
self.sync.state_information,
self.state_information,
)


Expand Down

Back | FazBrowse Home | New Git URL