add improved discovery (uses mdns to discover device types), use for …

…old behavior

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

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", []))
 

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

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()