From 7e2800d89aaabcb4b2ffa1fda7fe35f4bea65a85 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Sat, 22 Jun 2024 20:02:29 +0800
Subject: [PATCH] Make necessary changes for adopting the latest
 iphone_backup_decrypt

---
 Whatsapp_Chat_Exporter/__main__.py          |  9 +---
 Whatsapp_Chat_Exporter/ios_media_handler.py | 49 +++++++++------------
 2 files changed, 21 insertions(+), 37 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py
index a332181..fd6e2b6 100644
--- a/Whatsapp_Chat_Exporter/__main__.py
+++ b/Whatsapp_Chat_Exporter/__main__.py
@@ -188,13 +188,6 @@ def main():
         action='store_true',
         help="Use Whatsapp Business default files (iOS only)"
     )
-    parser.add_argument(
-        "--preserve-timestamp",
-        dest="preserve_timestamp",
-        default=False,
-        action='store_true',
-        help="Preserve the modification timestamp of the extracted files (iOS only)"
-    )
     parser.add_argument(
         "--wab",
         "--wa-backup",
@@ -403,7 +396,7 @@ def main():
             args.media = identifiers.DOMAIN
         if args.backup is not None:
             if not os.path.isdir(args.media):
-                ios_media_handler.extract_media(args.backup, identifiers, args.preserve_timestamp)
+                ios_media_handler.extract_media(args.backup, identifiers)
             else:
                 print("WhatsApp directory already exists, skipping WhatsApp file extraction.")
         if args.db is None:
diff --git a/Whatsapp_Chat_Exporter/ios_media_handler.py b/Whatsapp_Chat_Exporter/ios_media_handler.py
index 8fa0f65..6ece83d 100644
--- a/Whatsapp_Chat_Exporter/ios_media_handler.py
+++ b/Whatsapp_Chat_Exporter/ios_media_handler.py
@@ -3,55 +3,49 @@
 import shutil
 import sqlite3
 import os
-import time
 import getpass
-import threading
 from Whatsapp_Chat_Exporter.utility import WhatsAppIdentifier
 try:
     from iphone_backup_decrypt import EncryptedBackup, RelativePath
-    from iphone_backup_decrypt import FailedToDecryptError
 except ModuleNotFoundError:
     support_encrypted = False
 else:
     support_encrypted = True
 
 
-def extract_encrypted(base_dir, password, identifiers, bplist_reader=None):
+def extract_encrypted(base_dir, password, identifiers):
+    print("Trying to decrypt the iOS backup...", end="")
     backup = EncryptedBackup(backup_directory=base_dir, passphrase=password, cleanup=False, check_same_thread=False)
-    print("Decrypting WhatsApp database...", end="")
+    print("Done\nDecrypting WhatsApp database...", end="")
     try:
         backup.extract_file(
             relative_path=RelativePath.WHATSAPP_MESSAGES,
-            domain=identifiers.DOMAIN,
+            domain_like=identifiers.DOMAIN,
             output_filename=identifiers.MESSAGE
         )
         backup.extract_file(
             relative_path=RelativePath.WHATSAPP_CONTACTS,
-            domain=identifiers.DOMAIN,
+            domain_like=identifiers.DOMAIN,
             output_filename=identifiers.CONTACT
         )
-    except FailedToDecryptError:
+    except ValueError:
         print("Failed to decrypt backup: incorrect password?")
         exit()
     else:
         print("Done")
-    extract_thread = threading.Thread(
-        target=backup.extract_files_by_domain,
-        args=(identifiers.DOMAIN, identifiers.DOMAIN, bplist_reader)
+
+    def extract_progress_handler(file_id, domain, relative_path, n, total_files):
+        if n % 100 == 0:
+            print(f"Decrypting and extracting files...({n}/{total_files})", end="\r")   
+        return True
+
+    backup.extract_files(
+        domain_like=identifiers.DOMAIN,
+        output_folder=identifiers.DOMAIN,
+        preserve_folders=True,
+        filter_callback=extract_progress_handler
     )
-    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
+    print(f"All required files are decrypted and extracted.          ", end="\n")
     return backup
 
 
@@ -70,10 +64,7 @@ def is_encrypted(base_dir):
             return False
 
 
-def extract_media(base_dir, identifiers, preserve_timestamp=False):
-    if preserve_timestamp:
-        from Whatsapp_Chat_Exporter.bplist import BPListReader 
-        preserve_timestamp = BPListReader
+def extract_media(base_dir, identifiers):
     if is_encrypted(base_dir):
         if not support_encrypted:
             print("You don't have the dependencies to handle encrypted backup.")
@@ -82,7 +73,7 @@ def extract_media(base_dir, identifiers, preserve_timestamp=False):
             return False
         print("Encryption detected on the backup!")
         password = getpass.getpass("Enter the password for the backup:")
-        extract_encrypted(base_dir, password, identifiers, preserve_timestamp)
+        extract_encrypted(base_dir, password, identifiers)
     else:
         wts_db = os.path.join(base_dir, identifiers.MESSAGE[:2], identifiers.MESSAGE)
         contact_db = os.path.join(base_dir, identifiers.CONTACT[:2], identifiers.CONTACT)