Extend miio-extract-tokens to allow extraction from yeelight app db (#…

…462)

* Extend miio-extract-tokens to allow extraction from yeelight app database

* make hound happy

README.rst

@@ -27,11 +27,11 @@ Supported devices
 -  Xiaomi Mi Air Humidifier (:class:`miio.airhumidifier`)
 -  Xiaomi Mi Water Purifier (Basic support: Turn on & off) (:class:`miio.waterpurifier`)
 -  Xiaomi PM2.5 Air Quality Monitor (:class:`miio.airqualitymonitor`)
--  Xiaomi Smart WiFi Speaker (:class:`miio.wifispeaker`) (incomplete, please `feel free to help improve the support <https://github.com/rytilahti/python-miio/issues/69>`__)
+-  Xiaomi Smart WiFi Speaker (:class:`miio.wifispeaker`)
 -  Xiaomi Mi WiFi Repeater 2 (:class:`miio.wifirepeater`)
 -  Xiaomi Mi Smart Rice Cooker (:class:`miio.cooker`)
 -  Xiaomi Smartmi Fresh Air System (:class:`miio.airfresh`)
--  Yeelight light bulbs (:class:`miio.yeelight`) (only a very rudimentary support, use `python-yeelight <https://gitlab.com/stavros/python-yeelight/>`__ for a more complete support)
+-  :doc:`Yeelight light bulbs <yeelight>` (:class:`miio.yeelight`) (only a very rudimentary support, use `python-yeelight <https://gitlab.com/stavros/python-yeelight/>`__ for a more complete support)
 
 *Feel free to create a pull request to add support for new devices as
 well as additional features for supported devices.*

docs/index.rst

@@ -32,5 +32,6 @@ who have helped to extend this to cover not only the vacuum cleaner.
     plug
     ceil
     eyecare
+    yeelight
     API <miio>
     troubleshooting

docs/vacuum.rst

@@ -259,4 +259,4 @@ so it is also possible to pass dicts.
    :prog: mirobo
    :show-nested:
 
-:py:class:`API <miio.Vacuum>`
+:py:class:`API <miio.vacuum>`

miio/extract_tokens.py

@@ -1,8 +1,10 @@
 import logging
 import sqlite3
 import tempfile
+import json
 from pprint import pformat as pf
 from typing import Iterator
+import xml.etree.ElementTree as ET
 
 import attr
 import click
@@ -22,8 +24,31 @@ class DeviceConfig:
     ip = attr.ib()
     token = attr.ib()
     model = attr.ib()
+    everything = attr.ib(default=None)
 
 
+def read_android_yeelight(db) -> Iterator[DeviceConfig]:
+    """Read tokens from Yeelight's android backup."""
+    _LOGGER.info("Reading tokens from Yeelight Android DB")
+    xml = ET.parse(db)
+    devicelist = xml.find(".//set[@name='deviceList']")
+    if not devicelist:
+        _LOGGER.warning("Unable to find deviceList")
+        return []
+
+    for dev_elem in list(devicelist):
+        dev = json.loads(dev_elem.text)
+        ip = dev['localip']
+        mac = dev['mac']
+        model = dev['model']
+        name = dev['name']
+        token = dev['token']
+
+        config = DeviceConfig(name=name, ip=ip, mac=mac, model=model,
+                              token=token, everything=dev)
+
+        yield config
+
 class BackupDatabaseReader:
     """Main class for reading backup files.
     The main usage is following:
@@ -73,7 +98,8 @@ def read_apple(self) -> Iterator[DeviceConfig]:
             name = dev['ZNAME']
             token = BackupDatabaseReader.decrypt_ztoken(dev['ZTOKEN'])
 
-            config = DeviceConfig(name=name, mac=mac, ip=ip, model=model, token=token)
+            config = DeviceConfig(name=name, mac=mac, ip=ip, model=model,
+                                  token=token, everything=dev)
             yield config
 
     def read_android(self) -> Iterator[DeviceConfig]:
@@ -89,8 +115,8 @@ def read_android(self) -> Iterator[DeviceConfig]:
             name = dev['name']
             token = dev['token']
 
-            config = DeviceConfig(name=name, ip=ip, mac=mac,
-                                  model=model, token=token)
+            config = DeviceConfig(name=name, ip=ip, mac=mac, model=model,
+                                  token=token, everything=dev)
             yield config
 
     def read_tokens(self, db) -> Iterator[DeviceConfig]:
@@ -131,26 +157,46 @@ def main(backup, write_to_disk, password, dump_all, dump_raw):
      If the given file is an iOS backup, the tokens will be
      extracted (and decrypted if needed) automatically.
     """
-
+    def read_miio_database(tar):
+        DBFILE = "apps/com.xiaomi.smarthome/db/miio2.db"
+        try:
+            db = tar.extractfile(DBFILE)
+        except KeyError as ex:
+            click.echo("Unable to find miio database file %s: %s" % (
+                DBFILE, ex))
+            return []
+        if write_to_disk:
+            file = write_to_disk
+        else:
+            file = tempfile.NamedTemporaryFile()
+        with file as fp:
+            click.echo("Saving database to %s" % fp.name)
+            fp.write(db.read())
+
+            return list(reader.read_tokens(fp.name))
+
+    def read_yeelight_database(tar):
+        DBFILE = "apps/com.yeelight.cherry/sp/miot.xml"
+        _LOGGER.info("Trying to read %s", DBFILE)
+        try:
+            db = tar.extractfile(DBFILE)
+        except KeyError as ex:
+            click.echo("Unable to find yeelight database file %s: %s" % (
+                DBFILE, ex))
+            return []
+
+        return list(read_android_yeelight(db))
+
+    devices = []
     reader = BackupDatabaseReader(dump_raw)
     if backup.endswith(".ab"):
-        DBFILE = "apps/com.xiaomi.smarthome/db/miio2.db"
+
         with AndroidBackup(backup, stream=False) as f:
             tar = f.read_data(password)
-            try:
-                db = tar.extractfile(DBFILE)
-            except KeyError as ex:
-                click.echo("Unable to extract the database file %s: %s" % (DBFILE, ex))
-                return
-            if write_to_disk:
-                file = write_to_disk
-            else:
-                file = tempfile.NamedTemporaryFile()
-            with file as fp:
-                click.echo("Saving database to %s" % fp.name)
-                fp.write(db.read())
 
-                devices = list(reader.read_tokens(fp.name))
+            devices.extend(read_miio_database(tar))
+
+            devices.extend(read_yeelight_database(tar))
     else:
         devices = list(reader.read_tokens(backup))
 
@@ -162,6 +208,8 @@ def main(backup, write_to_disk, password, dump_all, dump_raw):
                        "\tToken: %s\n"
                        "\tMAC: %s" % (dev.name, dev.model,
                                       dev.ip, dev.token, dev.mac))
+            if dump_raw:
+                click.echo(dev)
 
 
 if __name__ == "__main__":