From 0777489010ef019f73ef362c16868d7afae6dd71 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Sun, 17 Sep 2017 21:24:51 +0200
Subject: [PATCH] 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: https://github.com/bluec0re/android-backup-tools/pull/1
---
 mirobo/extract_tokens.py | 73 +++++++++++++++++++++++++---------------
 1 file changed, 46 insertions(+), 27 deletions(-)

diff --git a/mirobo/extract_tokens.py b/mirobo/extract_tokens.py
index d848ddb85..4afa5eafb 100644
--- a/mirobo/extract_tokens.py
+++ b/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,18 +62,15 @@ 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(
@@ -72,9 +78,9 @@ def read_tokens(self, db):
             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!")
 
@@ -82,28 +88,41 @@ def read_tokens(self, db):
 @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__":