Implement firmware update functionality (#153)

Thanks for exploring the device in-depth goes to the dustcloud project,
which also describes how to build a custom firmware.

miio/device.py

@@ -5,6 +5,7 @@
 import construct
 import binascii
 from typing import Any, List, Optional  # noqa: F401
+from enum import Enum
 
 from .protocol import Message
 
@@ -21,6 +22,13 @@ class DeviceError(DeviceException):
     pass
 
 
+class UpdateState(Enum):
+    Downloading = "downloading"
+    Installing = "installing"
+    Failed = "failed"
+    Idle = "idle"
+
+
 class DeviceInfo:
     """Container of miIO device information.
     Hardware properties such as device model, MAC address, memory information,
@@ -288,6 +296,25 @@ def info(self) -> DeviceInfo:
         and harware and software versions."""
         return DeviceInfo(self.send("miIO.info", []))
 
+    def update(self, url: str, md5: str):
+        """Start an OTA update."""
+        payload = {
+            "mode": "normal",
+            "install": "1",
+            "app_url": url,
+            "file_md5": md5,
+            "proc": "dnld install"
+        }
+        return self.send("miIO.ota", payload)[0] == "ok"
+
+    def update_progress(self) -> int:
+        """Return current update progress [0-100]."""
+        return self.send("miIO.get_ota_progress", [])[0]
+
+    def update_state(self):
+        """Return current update state."""
+        return UpdateState(self.send("miIO.get_ota_state", [])[0])
+
     @property
     def _id(self) -> int:
         """Increment and return the sequence id."""

miio/updater.py

@@ -0,0 +1,91 @@
+from http.server import HTTPServer, BaseHTTPRequestHandler
+import hashlib
+import logging
+import netifaces
+from os.path import basename
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class SingleFileHandler(BaseHTTPRequestHandler):
+    """A simplified handler just returning the contents of a buffer."""
+    def __init__(self, request, client_address, server):
+        self.payload = server.payload
+        self.server = server
+
+        super().__init__(request, client_address, server)
+
+    def handle_one_request(self):
+        self.server.got_request = True
+        self.raw_requestline = self.rfile.readline()
+
+        if not self.parse_request():
+            _LOGGER.error("unable to parse request: %s" % self.raw_requestline)
+            return
+
+        self.send_response(200)
+        self.send_header('Content-type', 'application/octet-stream')
+        self.send_header('Content-Length', len(self.payload))
+        self.end_headers()
+        self.wfile.write(self.payload)
+
+
+class OneShotServer:
+    """A simple HTTP server for serving an update file.
+
+    The server will be started in an emphemeral port, and will only accept
+    a single request to keep it simple."""
+    def __init__(self, file, interface=None):
+        addr = ('', 0)
+        self.server = HTTPServer(addr, SingleFileHandler)
+        setattr(self.server, "got_request", False)
+
+        self.addr, self.port = self.server.server_address
+        self.server.timeout = 10
+
+        _LOGGER.info("Serving on %s:%s, timeout %s" % (self.addr, self.port,
+                                                       self.server.timeout))
+
+        self.file = basename(file)
+        with open(file, 'rb') as f:
+            self.payload = f.read()
+            self.server.payload = self.payload
+            self.md5 = hashlib.md5(self.payload).hexdigest()
+            _LOGGER.info("Using local %s (md5: %s)" % (file, self.md5))
+
+    @staticmethod
+    def find_local_ip():
+        ifaces_without_lo = [x for x in netifaces.interfaces()
+                             if not x.startswith("lo")]
+        _LOGGER.debug("available interfaces: %s" % ifaces_without_lo)
+
+        for iface in ifaces_without_lo:
+            addresses = netifaces.ifaddresses(iface)
+            if netifaces.AF_INET not in addresses:
+                _LOGGER.debug("%s has no ipv4 addresses, skipping" % iface)
+                continue
+            for entry in addresses[netifaces.AF_INET]:
+                _LOGGER.debug("Got addr: %s" % entry['addr'])
+                return entry['addr']
+
+    def url(self, ip=None):
+        if ip is None:
+            ip = OneShotServer.find_local_ip()
+
+        url = "http://%s:%s/%s" % (ip, self.port, self.file)
+        return url
+
+    def serve_once(self):
+        self.server.handle_request()
+        if getattr(self.server, "got_request"):
+            _LOGGER.info("Got a request, shold be downloading now.")
+            return True
+        else:
+            _LOGGER.error("No request was made..")
+            return False
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    upd = OneShotServer("/tmp/test")
+    upd.serve_once()

miio/vacuum_cli.py

@@ -7,11 +7,15 @@
 import json
 import time
 import pathlib
+import threading
+from tqdm import tqdm
 from appdirs import user_cache_dir
 from pprint import pformat as pf
 from typing import Any  # noqa: F401
 from miio.click_common import (ExceptionHandlerGroup, validate_ip,
                                validate_token)
+from .device import UpdateState
+from .updater import OneShotServer
 import miio  # noqa: E402
 
 _LOGGER = logging.getLogger(__name__)
@@ -425,26 +429,46 @@ def sound(vac: miio.Vacuum, volume: int, test_mode: bool):
 
 @cli.command()
 @click.argument('url')
-@click.argument('md5sum')
-@click.argument('sid', type=int)
+@click.argument('md5sum', required=False, default=None)
+@click.argument('sid', type=int, required=False, default=10000)
 @pass_dev
 def install_sound(vac: miio.Vacuum, url: str, md5sum: str, sid: int):
     """Install a sound."""
     click.echo("Installing from %s (md5: %s) for id %s" % (url, md5sum, sid))
-    click.echo(vac.install_sound(url, md5sum, sid))
+
+    local_url = None
+    server = None
+    if url.startswith("http"):
+        if md5sum is None:
+            click.echo("You need to pass md5 when using URL for updating.")
+            return
+        local_url = url
+    else:
+        server = OneShotServer(url)
+        local_url = server.url()
+        md5sum = server.md5
+
+        t = threading.Thread(target=server.serve_once)
+        t.start()
+        click.echo("Hosting file at %s" % local_url)
+
+    click.echo(vac.install_sound(local_url, md5sum, sid))
 
     progress = vac.sound_install_progress()
     while progress.is_installing:
-        print(progress)
         progress = vac.sound_install_progress()
-        time.sleep(0.1)
+        print("%s (%s %%)" % (progress.state.name, progress.progress))
+        time.sleep(1)
 
     progress = vac.sound_install_progress()
 
-    if progress.progress == 100 and progress.error == 0:
-        click.echo("Installation of sid '%s' complete!" % progress.sid)
-    else:
+    if progress.is_errored:
         click.echo("Error during installation: %s" % progress.error)
+    else:
+        click.echo("Installation of sid '%s' complete!" % sid)
+
+    if server is not None:
+        t.join()
 
 @cli.command()
 @pass_dev
@@ -481,6 +505,69 @@ def configure_wifi(vac: miio.Vacuum, ssid: str, password: str,
     click.echo(vac.configure_wifi(ssid, password, uid, timezone))
 
 
+@cli.command()
+@pass_dev
+def update_status(vac: miio.Vacuum):
+    """Return update state and progress."""
+    update_state = vac.update_state()
+    click.echo("Update state: %s" % update_state)
+
+    if update_state == UpdateState.Downloading:
+        click.echo("Update progress: %s" % vac.update_progress())
+
+
+@cli.command()
+@click.argument('url', required=True)
+@click.argument('md5', required=False, default=None)
+@pass_dev
+def update_firmware(vac: miio.Vacuum, url: str, md5: str):
+    """Update device firmware.
+
+    If `file` starts with http* it is expected to be an URL.
+     In that case md5sum of the file has to be given."""
+
+    # TODO Check that the device is in updateable state.
+
+    click.echo("Going to update from %s" % url)
+    if url.lower().startswith("http"):
+        if md5 is None:
+            click.echo("You need to pass md5 when using URL for updating.")
+            return
+
+        click.echo("Using %s (md5: %s)" % (url, md5))
+    else:
+        server = OneShotServer(url)
+        url = server.url()
+
+        t = threading.Thread(target=server.serve_once)
+        t.start()
+        click.echo("Hosting file at %s" % url)
+        md5 = server.md5
+
+    update_res = vac.update(url, md5)
+    if update_res:
+        click.echo("Update started!")
+    else:
+        click.echo("Starting the update failed: %s" % update_res)
+
+    with tqdm(total=100) as t:
+        state = vac.update_state()
+        while state == UpdateState.Downloading:
+            try:
+                state = vac.update_state()
+                progress = vac.update_progress()
+            except:  # we may not get our messages through during upload
+                continue
+
+            if state == UpdateState.Installing:
+                click.echo("Installation started, please wait until the vacuum reboots")
+                break
+
+            t.update(progress - t.n)
+            t.set_description("%s" % state.name)
+            time.sleep(1)
+
+
 @cli.command()
 @click.argument('cmd', required=True)
 @click.argument('parameters', required=False)

miio/vacuumcontainers.py

@@ -522,7 +522,13 @@ def error(self) -> int:
     @property
     def is_installing(self) -> bool:
         """True if install is in progress."""
-        return self.sid != 0 and self.progress < 100 and self.error == 0
+        return (self.state == SoundInstallState.Downloading or
+                self.state == SoundInstallState.Installing)
+
+    @property
+    def is_errored(self) -> bool:
+        """True if the state has an error, use `error`to access it."""
+        return self.state == SoundInstallState.Error
 
     def __repr__(self) -> str:
         return "<SoundInstallStatus sid: %s (state: %s, error: %s)" \