Update the way to handle encrypted iOS backup file

Since this commit, iphone_backup_decrypt must be re-installed

Whatsapp_Chat_Exporter/extract_iphone_media.py

@@ -3,70 +3,57 @@
 import shutil
 import sqlite3
 import os
+import time
 import getpass
+import threading
 try:
-    from iphone_backup_decrypt import EncryptedBackup, RelativePath
+    from iphone_backup_decrypt import EncryptedBackup, RelativePath, FailedToDecryptError, Domain
 except ModuleNotFoundError:
     support_encrypted = False
 else:
     support_encrypted = True
 
 
 def extract_encrypted(base_dir, password):
-    backup = EncryptedBackup(backup_directory=base_dir, passphrase=password)
+    backup = EncryptedBackup(backup_directory=base_dir, passphrase=password, cleanup=False, check_same_thread=False)
     print("Decrypting WhatsApp database...")
-    backup.extract_file(relative_path=RelativePath.WHATSAPP_MESSAGES,
+    try:
+        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"Extracting 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"Extracting media...({i}/{total_row_number})", end="\r")
-    print(f"Extracting media...({total_row_number}/{total_row_number})", end="\n")
+        backup.extract_file(relative_path=RelativePath.WHATSAPP_CONTACTS,
+                            output_filename="b8548dc30aa1030df0ce18ef08b882cf7ab5212f")
+    except FailedToDecryptError:
+        print("Failed to decrypt backup: incorrect password?")
+        exit()
+    extract_thread = threading.Thread(
+        target=backup.extract_files_by_domain,
+        args=(Domain.WHATSAPP, Domain.WHATSAPP)
+    )
+    extract_thread.daemon = True
+    extract_thread.start()
+    dot = 0
+    while extract_thread.is_alive():
+        print(f"Decrypting and extracting files{'.' * dot}{' ' * (3 - dot)}", end="\r")
+        if dot < 3:
+            dot += 1
+            time.sleep(0.5)
+        else:
+            dot = 0
+            time.sleep(0.4)
+    print(f"All required files decrypted and extracted.", end="\n")
+    extract_thread.handled = True
+    return backup
 
 
 def is_encrypted(base_dir):
-    with sqlite3.connect(f"{base_dir}/Manifest.db") as f:
+    with sqlite3.connect(os.path.join(base_dir, "Manifest.db")) as f:
         c = f.cursor()
         try:
             c.execute("""SELECT count()
                     FROM Files
                     """)
+        except sqlite3.OperationalError as e:
+            raise e  # These error cannot be used to determine if the backup is encrypted
         except sqlite3.DatabaseError:
             return True
         else:
@@ -80,7 +67,8 @@ def extract_media(base_dir):
             print("Read more on how to deal with encrypted backup:")
             print("https://github.com/KnugiHK/Whatsapp-Chat-Exporter/blob/main/README.md#usage")
             return False
-        password = getpass.getpass("Enter the password:")
+        print("Encryption detected on the backup!")
+        password = getpass.getpass("Enter the password for the backup:")
         extract_encrypted(base_dir, password)
     else:
         wts_db = os.path.join(base_dir, "7c/7c7fba66680ef796b916b067077cc246adacf01d")
@@ -94,15 +82,15 @@ def extract_media(base_dir):
             print("Contact database not found.")
             exit()
         else:
-            shutil.copyfile(contact_db, "b8548dc30aa1030df0ce18ef08b882cf7ab5212f")    
+            shutil.copyfile(contact_db, "b8548dc30aa1030df0ce18ef08b882cf7ab5212f")
         _wts_id = "AppDomainGroup-group.net.whatsapp.WhatsApp.shared"
         with sqlite3.connect(os.path.join(base_dir, "Manifest.db")) as manifest:
             manifest.row_factory = sqlite3.Row
             c = manifest.cursor()
             c.execute(
                 f"""SELECT count()
-                   FROM Files
-                   WHERE domain = '{_wts_id}'"""
+                    FROM Files
+                    WHERE domain = '{_wts_id}'"""
             )
             total_row_number = c.fetchone()[0]
             print(f"Extracting WhatsApp files...(0/{total_row_number})", end="\r")