From a4a10a08ff65c128eceb3485bdcdd4b71a85edb1 Mon Sep 17 00:00:00 2001 From: Teemu Rytilahti Date: Tue, 22 Aug 2017 01:58:01 +0200 Subject: [PATCH] add improved discovery (uses mdns to discover device types), use for old behavior --- mirobo/__init__.py | 1 + mirobo/device.py | 21 ++++++++++----- mirobo/discovery.py | 63 ++++++++++++++++++++++++++++++++++++++++++++ mirobo/vacuum_cli.py | 16 +++++++---- 4 files changed, 90 insertions(+), 11 deletions(-) create mode 100644 mirobo/discovery.py diff --git a/mirobo/__init__.py b/mirobo/__init__.py index fe83a0da6..09ae830e0 100644 --- a/mirobo/__init__.py +++ b/mirobo/__init__.py @@ -9,3 +9,4 @@ from mirobo.ceil import Ceil from mirobo.philips_eyecare import PhilipsEyecare from mirobo.device import Device, DeviceException +from mirobo.discovery import Discovery diff --git a/mirobo/device.py b/mirobo/device.py index a8f44021f..8dff9e79f 100644 --- a/mirobo/device.py +++ b/mirobo/device.py @@ -37,11 +37,14 @@ def raw(self): class Device: - def __init__(self, ip: str, token: str, + def __init__(self, ip: str = None, token: str = None, start_id: int=0, debug: int=0) -> None: self.ip = ip self.port = 54321 - self.token = bytes.fromhex(token) + if token is None: + token = 32 * '0' + if token is not None: + self.token = bytes.fromhex(token) self.debug = debug self._timeout = 5 @@ -59,9 +62,11 @@ def do_discover(self): self._device_ts = m.header.value.ts if self.debug > 1: _LOGGER.debug(m) - _LOGGER.debug("Discovered %s %s with ts: %s" % (self._devtype, - self._serial, - self._device_ts)) + _LOGGER.debug("Discovered %s %s with ts: %s, token: %s" % ( + self._devtype, + self._serial, + self._device_ts, + codecs.encode(m.checksum, 'hex'))) else: _LOGGER.error("Unable to discover a device at address %s", self.ip) raise DeviceException("Unable to discover the device %s" % self.ip) @@ -92,7 +97,7 @@ def discover(addr: str=None) -> Any: try: data, addr = s.recvfrom(1024) m = Message.parse(data) # type: Message - # _LOGGER.debug("Got a response: %s" % m) + _LOGGER.debug("Got a response: %s" % m) if not is_broadcast: return m @@ -167,6 +172,10 @@ def send(self, command: str, parameters: Any=None, retry_count=3) -> Any: return self.send(command, parameters, retry_count-1) raise DeviceException from ex + def raw_command(self, cmd, params): + """Send a raw command to the robot.""" + return self.send(cmd, params) + def info(self): return DeviceInfo(self.send("miIO.info", [])) diff --git a/mirobo/discovery.py b/mirobo/discovery.py new file mode 100644 index 000000000..dd1c5bd08 --- /dev/null +++ b/mirobo/discovery.py @@ -0,0 +1,63 @@ +import logging +import zeroconf +import ipaddress +import inspect +import codecs +from typing import Any, List, Union, Callable +from mirobo import Message, Device, Vacuum + +_LOGGER = logging.getLogger(__name__) + + +def other_package_info(info, desc): + return "%s @ %s, check %s" % (info.name, ipaddress.ip_address(info.address), desc) + + +class Listener: + def __init__(self): + self.found_devices = {} + def _check_if_supported(self, info, addr): + name = info.name + for k, v in Discovery._mdns_device_map.items(): + if name.startswith(k): + if inspect.isclass(v): + dev = v(ip=addr) + m = dev.do_discover() + dev.token = m.checksum + _LOGGER.info("Found supported '%s' at %s:%s (%s) token: %s" % ( + v.__name__, addr, info.port, name, dev.token)) + return dev + elif callable(v): + _LOGGER.info(v(info)) + dev = Device(ip=addr) + _LOGGER.info("token: %s" % codecs.encode(dev.do_discover().checksum, 'hex')) + return None + _LOGGER.warning("Found unsupported device %s at %s, please report to developers" % (name, addr)) + return None + + def add_service(self, zeroconf, type, name): + info = zeroconf.get_service_info(type, name) + addr = str(ipaddress.ip_address(info.address)) + if addr not in self.found_devices: + dev = self._check_if_supported(info, addr) + self.found_devices[addr] = dev + + +class Discovery: + _mdns_device_map = { + "rockrobo-vacuum-v1": Vacuum, + "yeelink-light-": lambda x: other_package_info(x, "python-yeelight package"), + "lumi-gateway-": lambda x: other_package_info(x, "https://github.com/Danielhiversen/PyXiaomiGateway") + } # type: Dict[str, Union[Callable, Device]] + + @staticmethod + def discover_mdns(): + _LOGGER.info("Discovering devices with mDNS, press any key to quit...") + + listener = Listener() + browser = zeroconf.ServiceBrowser(zeroconf.Zeroconf(), "_miio._udp.local.", listener) + + input() # to keep execution running until a key is pressed + browser.cancel() + + return listener.found_devices \ No newline at end of file diff --git a/mirobo/vacuum_cli.py b/mirobo/vacuum_cli.py index 244455c94..f7c8df65f 100644 --- a/mirobo/vacuum_cli.py +++ b/mirobo/vacuum_cli.py @@ -17,7 +17,7 @@ import mirobo # noqa: E402 _LOGGER = logging.getLogger(__name__) -pass_dev = click.make_pass_decorator(mirobo.Vacuum) +pass_dev = click.make_pass_decorator(mirobo.Device, ensure=True) def validate_ip(ctx, param, value): @@ -57,6 +57,7 @@ def cli(ctx, ip: str, token: str, debug: int, id_file: str): # if we are scanning, we do not try to connect. if ctx.invoked_subcommand == "discover": + ctx.obj = "discover" return if ip is None or token is None: @@ -75,6 +76,7 @@ def cli(ctx, ip: str, token: str, debug: int, id_file: str): pass vac = mirobo.Vacuum(ip, token, start_id, debug) + vac.manual_seqnum = manual_seq _LOGGER.debug("Connecting to %s with token %s", ip, token) @@ -88,19 +90,23 @@ def cli(ctx, ip: str, token: str, debug: int, id_file: str): @cli.resultcallback() @pass_dev def cleanup(vac: mirobo.Vacuum, **kwargs): + if vac.ip is None: # dummy Device for discovery, skip teardown + return id_file = kwargs['id_file'] seqs = {'seq': vac.raw_id, 'manual_seq': vac.manual_seqnum} _LOGGER.debug("Writing %s to %s" % (seqs, id_file)) with open(id_file, 'w') as f: json.dump(seqs, f) - #with open(id_file, 'w') as f: - # f.write(str(vac.raw_id)) @cli.command() -def discover(): +@click.option('--handshake', type=bool, default=False) +def discover(handshake): """Search for robots in the network.""" - mirobo.Vacuum.discover() + if handshake: + mirobo.Vacuum.discover() + else: + mirobo.Discovery.discover_mdns() @cli.command()