From be7e20317da7c7f06e31e9be2e847fae3886def6 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Mon, 10 May 2021 13:33:13 +0800
Subject: [PATCH] Add support for encrypted iPhone backup

---
 extract_iphone_media.py | 138 ++++++++++++++++++++++++++++++----------
 1 file changed, 105 insertions(+), 33 deletions(-)

diff --git a/extract_iphone_media.py b/extract_iphone_media.py
index 38d593f..3646b2a 100644
--- a/extract_iphone_media.py
+++ b/extract_iphone_media.py
@@ -3,43 +3,115 @@
 import shutil
 import sqlite3
 import os
+import getpass
+try:
+    from iphone_backup_decrypt import EncryptedBackup, RelativePath
+except:
+    support_encrypted = False
+else:
+    support_encrypted = True
 
+def extract_encrypted(base_dir, password):
+    backup = EncryptedBackup(backup_directory=base_dir, passphrase=password)
+    print("Decrypting WhatsApp database...")
+    backup.extract_file(relative_path=RelativePath.WHATSAPP_MESSAGES, output_filename="7c7fba66680ef796b916b067077cc246adacf01d")
+    backup.extract_file(relative_path=RelativePath.WHATSAPP_CONTACTS, output_filename="ContactsV2.sqlite")
+    data = backup.execute_sql("""SELECT count()
+                                FROM Files
+                                WHERE relativePath
+                                    LIKE 'Message/Media/%'"""
+    )
+    total_row_number = data[0][0]
+    print(f"Gathering media...(0/{total_row_number})", end="\r")
+    data = backup.execute_sql("""SELECT fileID,
+                                        relativePath,
+                                        flags,
+                                        file
+                                FROM Files
+                                WHERE relativePath
+                                    LIKE 'Message/Media/%'"""
+    )
+    if not os.path.isdir("Message"):
+        os.mkdir("Message")
+    if not os.path.isdir("Message/Media"):
+        os.mkdir("Message/Media")
+    i = 0
+    for row in data:
+        destination = row[1]
+        hashes = row[0]
+        folder = hashes[:2]
+        flags = row[2]
+        file = row[3]
+        if flags == 2:
+            try: 
+                os.mkdir(destination)
+            except FileExistsError:
+                pass
+        elif flags == 1:
+            decrypted = backup.decrypt_inner_file(file_id=hashes, file_bplist=file)
+            with open(destination, "wb") as f:
+                f.write(decrypted)
+        i += 1
+        if i % 100 == 0:
+            print(f"Gathering media...({i}/{total_row_number})", end="\r")
+    print(f"Gathering media...({total_row_number}/{total_row_number})", end="\r")
+
+def is_encrypted(base_dir):
+    with sqlite3.connect(f"{base_dir}/Manifest.db") as f:
+        c = f.cursor()
+        try:
+            c.execute("""SELECT count()
+                    FROM Files
+                    """)
+        except sqlite3.DatabaseError:
+            return True
+        else:
+            return False
 
 def extract_media(base_dir):
-    with sqlite3.connect(f"{base_dir}/Manifest.db") as manifest:
-        c = manifest.cursor()
-        c.execute("""SELECT count()
-                    FROM Files
-                    WHERE relativePath
-                        LIKE 'Message/Media/%'""")
-        total_row_number = c.fetchone()[0]
-        print(f"Gathering media...(0/{total_row_number})", end="\r")
-        c.execute("""SELECT fileID,
-                            relativePath,
-                            flags
-                    FROM Files
-                    WHERE relativePath
-                        LIKE 'Message/Media/%'""")
-        row = c.fetchone()
-        if not os.path.isdir("Message"):
-            os.mkdir("Message")
-        if not os.path.isdir("Message/Media"):
-            os.mkdir("Message/Media")
-        i = 0
-        while row is not None:
-            destination = row[1]
-            hashes = row[0]
-            folder = hashes[:2]
-            flags = row[2]
-            if flags == 2:
-                os.mkdir(destination)
-            elif flags == 1:
-                shutil.copyfile(f"{base_dir}/{folder}/{hashes}", destination)
-            i += 1
-            if i % 100 == 0:
-                print(f"Gathering media...({i}/{total_row_number})", end="\r")
+    if is_encrypted(base_dir):
+        if not support_encrypted:
+            print("You don't have the dependencies to handle encrypted backup.")
+            print("Read more about how to deal with encrypted backup:")
+            print("https://github.com/KnugiHK/Whatsapp-Chat-Exporter/blob/main/README.md#encrypted-iphone-backup")
+            return False
+        password = getpass.getpass("Enter the password:")
+        extract_encrypted(base_dir, password)
+    else:
+        with sqlite3.connect(f"{base_dir}/Manifest.db") as manifest:
+            c = manifest.cursor()
+            c.execute("""SELECT count()
+                        FROM Files
+                        WHERE relativePath
+                            LIKE 'Message/Media/%'""")
+            total_row_number = c.fetchone()[0]
+            print(f"Gathering media...(0/{total_row_number})", end="\r")
+            c.execute("""SELECT fileID,
+                                relativePath,
+                                flags
+                        FROM Files
+                        WHERE relativePath
+                            LIKE 'Message/Media/%'""")
             row = c.fetchone()
-        print(f"Gathering media...({total_row_number}/{total_row_number})", end="\r")
+            if not os.path.isdir("Message"):
+                os.mkdir("Message")
+            if not os.path.isdir("Message/Media"):
+                os.mkdir("Message/Media")
+            i = 0
+            while row is not None:
+                destination = row[1]
+                hashes = row[0]
+                folder = hashes[:2]
+                flags = row[2]
+                if flags == 2:
+                    os.mkdir(destination)
+                elif flags == 1:
+                    shutil.copyfile(f"{base_dir}/{folder}/{hashes}", destination)
+                i += 1
+                if i % 100 == 0:
+                    print(f"Gathering media...({i}/{total_row_number})", end="\r")
+                row = c.fetchone()
+            print(f"Gathering media...({total_row_number}/{total_row_number})", end="\r")
 
 
 if __name__ == "__main__":