python-kasa/pyHS100/discover.py

import socket
import logging
import json
from typing import Dict, Type

from pyHS100 import TPLinkSmartHomeProtocol, SmartDevice, SmartPlug, SmartBulb

_LOGGER = logging.getLogger(__name__)


class Discover:
    DISCOVERY_QUERY = {"system": {"get_sysinfo": None},
                       "emeter": {"get_realtime": None}}

    @staticmethod
    def discover(protocol: TPLinkSmartHomeProtocol = None,
                 port: int = 9999,
                 timeout: int = 3) -> Dict[str, SmartDevice]:
        """
        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
        :param timeout: How long to wait for responses, defaults to 5
        :param port: port to send broadcast messages, defaults to 9999.
        :rtype: dict
        :return: Array of json objects {"ip", "port", "sys_info"}
        """
        if protocol is None:
            protocol = TPLinkSmartHomeProtocol()

        target = "255.255.255.255"

        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)

        encrypted_req = protocol.encrypt(req)
        sock.sendto(encrypted_req[4:], (target, port))

        devices = {}
        _LOGGER.debug("Waiting %s seconds for responses...", timeout)

        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 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)
        return devices

    @staticmethod
    def discover_single(ip_address: str,
                        protocol: TPLinkSmartHomeProtocol = None
                        ) -> SmartDevice:
        """
        Similar to discover(), except only return device object for single
        passed device given by IP address.

        :param ip_address: IP address 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()

        info = protocol.query(ip_address, Discover.DISCOVERY_QUERY)

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

    @staticmethod
    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"]
            if "type" in sysinfo:
                type = sysinfo["type"]
            elif "mic_type" in sysinfo:
                type = sysinfo["mic_type"]
            else:
                _LOGGER.error("Unable to find the device type field!")
                type = "UNKNOWN"
        else:
            _LOGGER.error("No 'system' nor 'get_sysinfo' in response")
        if "smartplug" in type.lower():
            return SmartPlug
        elif "smartbulb" in type.lower():
            return SmartBulb

        return None