Extract Android backups, yield devices instead of just echoing

This moves the functionality towards automatic config generation out of backups.
The extraction depends on code which is not available in pypi, and requires some
changes to be useful for us: bluec0re/android-backup-tools#1

mirobo/extract_tokens.py

@@ -4,11 +4,19 @@
 import sqlite3
 from Crypto.Cipher import AES
 from pprint import pformat as pf
+import attr
+
+@attr.s
+class DeviceConfig:
+    name = attr.ib()
+    mac = attr.ib()
+    ip = attr.ib()
+    token = attr.ib()
+    model = attr.ib()
 
 
 class BackupDatabaseReader:
-    def __init__(self, dump_all=False, dump_raw=False):
-        self.dump_all = dump_all
+    def __init__(self, dump_raw=False):
         self.dump_raw = dump_raw
 
     @staticmethod
@@ -39,8 +47,9 @@ def read_apple(self):
             model = dev['ZMODEL']
             name = dev['ZNAME']
             token = BackupDatabaseReader.decrypt_ztoken(dev['ZTOKEN'])
-            if ip or self.dump_all:
-                click.echo("%s\n\tModel: %s\n\tIP address: %s\n\tToken: %s\n\tMAC: %s" % (name, model, ip, token, mac))
+
+            config = DeviceConfig(name=name, mac=mac, ip=ip, model=model, token=token)
+            yield config
 
     def read_android(self):
         click.echo("Reading tokens from Android DB")
@@ -53,57 +62,67 @@ def read_android(self):
             model = dev['model']
             name = dev['name']
             token = dev['token']
-            if ip or self.dump_all:
-                click.echo("%s\n\tModel: %s\n\tIP address: %s\n\tToken: %s\n\tMAC: %s" % (name, model, ip, token, mac))
 
-    def dump_to_file(self, fp):
-        fp.open()
-        self.db.seek(0)  # go to the beginning
-        click.echo("Saving db to %s" % fp)
-        fp.write(self.db.read())
+            config = DeviceConfig(name=name, ip=ip, mac=mac, model=model, token=token)
+            yield config
 
     def read_tokens(self, db):
         self.db = db
+        print("reading database from %s" % db)
         self.conn = sqlite3.connect(db)
+
         self.conn.row_factory = sqlite3.Row
         with self.conn:
             is_android = self.conn.execute(
                 "SELECT name FROM sqlite_master WHERE type='table' AND name='devicerecord';").fetchone() is not None
             is_apple = self.conn.execute(
                 "SELECT name FROM sqlite_master WHERE type='table' AND name='ZDEVICE'").fetchone() is not None
             if is_android:
-                self.read_android()
+                yield from self.read_android()
             elif is_apple:
-                self.read_apple()
+                yield from self.read_apple()
             else:
                 click.echo("Error, unknown database type!")
 
 
 @click.command()
 @click.argument('backup')
 @click.option('--write-to-disk', type=click.File('wb'), help='writes sqlite3 db to a file for debugging')
+@click.option('--password', type=str, help='password if the android database is encrypted')
 @click.option('--dump-all', is_flag=True, default=False, help='dump devices without ip addresses')
 @click.option('--dump-raw', is_flag=True, help='dumps raw rows')
-def main(backup, write_to_disk, dump_all, dump_raw):
+def main(backup, write_to_disk, password, dump_all, dump_raw):
     """Reads device information out from an sqlite3 DB.
-     If the given file is a .tar file, the file will be extracted
-     and the database automatically located (out of Android backups).
+     If the given file is an .ab file, it will be extracted
+     and the database automatically located (Android backups).
     """
-    reader = BackupDatabaseReader(dump_all, dump_raw)
-    if backup.endswith(".tar"):
+
+    reader = BackupDatabaseReader(dump_raw)
+    from android_backup import AndroidBackup
+    if backup.endswith(".ab"):
         DBFILE = "apps/com.xiaomi.smarthome/db/miio2.db"
-        with tarfile.open(backup) as f:
-            click.echo("Opened %s" % backup)
-            db = f.extractfile(DBFILE)
-            with tempfile.NamedTemporaryFile() as fp:
-                click.echo("Extracting to %s" % fp.name)
+        with AndroidBackup(backup) as f:
+            tar = f.read_data(password)
+            try:
+                db = tar.extractfile(DBFILE)
+            except KeyError as ex:
+                click.echo("Unable to extract the device 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())
-                if write_to_disk:
-                    reader.dump_to_file(write_to_disk)
 
-                reader.read_tokens(fp.name)
+                devices = list(reader.read_tokens(fp.name))
     else:
-        reader.read_tokens(backup)
+        devices = list(reader.read_tokens(backup))
+
+    for dev in devices:
+        if dev.ip or dump_all:
+            click.echo("%s\n\tModel: %s\n\tIP address: %s\n\tToken: %s\n\tMAC: %s" % (dev.name, dev.model, dev.ip, dev.token, dev.mac))
 
 
 if __name__ == "__main__":