Restore concurrent backup decryption

src/mvt/common/help.py

@@ -28,6 +28,7 @@ HELP_MSG_DECRYPT_BACKUP = "Decrypt an encrypted iTunes backup"
 HELP_MSG_BACKUP_DESTINATION = (
     "Path to the folder where the decrypted backup should be stored"
 )
+HELP_MSG_DECRYPT_JOBS = "Number of files to decrypt concurrently"
 HELP_MSG_IOS_BACKUP_PASSWORD = (
     "Password to use to decrypt the backup (or, set the {MVT_IOS_BACKUP_PASSWORD} "
     "environment variable)"

src/mvt/ios/cli.py

@@ -33,6 +33,7 @@ from mvt.common.help import (
     HELP_MSG_VERSION,
     HELP_MSG_DECRYPT_BACKUP,
     HELP_MSG_BACKUP_DESTINATION,
+    HELP_MSG_DECRYPT_JOBS,
     HELP_MSG_IOS_BACKUP_PASSWORD,
     HELP_MSG_BACKUP_KEYFILE,
     HELP_MSG_HASHES,
@@ -58,7 +59,11 @@ from mvt.common.password import prompt_password
 from .cmd_check_backup import CmdIOSCheckBackup
 from .cmd_check_fs import CmdIOSCheckFS
 from .cmd_check_sysdiagnose import CmdIOSCheckSysdiagnose
-from .decrypt import DecryptBackup
+from .decrypt import (
+    DEFAULT_DECRYPT_WORKERS,
+    MAX_DECRYPT_WORKERS,
+    DecryptBackup,
+)
 from .modules.backup import BACKUP_MODULES
 from .modules.fs import FS_MODULES
 from .modules.mixed import MIXED_MODULES
@@ -162,6 +167,13 @@ def completion(ctx, shell, install):
     "decrypt-backup", context_settings=CONTEXT_SETTINGS, help=HELP_MSG_DECRYPT_BACKUP
 )
 @click.option("--destination", "-d", required=True, help=HELP_MSG_BACKUP_DESTINATION)
+@click.option(
+    "--jobs",
+    type=click.IntRange(1, MAX_DECRYPT_WORKERS),
+    default=DEFAULT_DECRYPT_WORKERS,
+    show_default=True,
+    help=HELP_MSG_DECRYPT_JOBS,
+)
 @click.option(
     "--password",
     "-p",
@@ -180,8 +192,8 @@ def completion(ctx, shell, install):
 @click.option("--hashes", "-H", is_flag=True, help=HELP_MSG_HASHES)
 @click.argument("BACKUP_PATH", type=click.Path(exists=True))
 @click.pass_context
-def decrypt_backup(ctx, destination, password, key_file, hashes, backup_path):
-    backup = DecryptBackup(backup_path, destination)
+def decrypt_backup(ctx, destination, jobs, password, key_file, hashes, backup_path):
+    backup = DecryptBackup(backup_path, destination, max_workers=jobs)
 
     if key_file:
         if MVT_IOS_BACKUP_PASSWORD in os.environ:

src/mvt/ios/decrypt.py

@@ -12,6 +12,13 @@ import plistlib
 import shutil
 import sqlite3
 import tempfile
+from concurrent.futures import (
+    ALL_COMPLETED,
+    FIRST_COMPLETED,
+    Future,
+    ThreadPoolExecutor,
+    wait,
+)
 from pathlib import Path
 from typing import Optional
 
@@ -21,6 +28,9 @@ from iphone_backup_decrypt.utils import FilePlist
 
 log = logging.getLogger(__name__)
 
+DEFAULT_DECRYPT_WORKERS = 4
+MAX_DECRYPT_WORKERS = 32
+
 # Import pbkdf2_hmac from the same source iphone_backup_decrypt uses internally,
 # so our key derivation is consistent with theirs.
 try:
@@ -178,19 +188,68 @@ class DecryptBackup:
 
     """
 
-    def __init__(self, backup_path: str, dest_path: Optional[str] = None) -> None:
+    def __init__(
+        self,
+        backup_path: str,
+        dest_path: Optional[str] = None,
+        max_workers: int = DEFAULT_DECRYPT_WORKERS,
+    ) -> None:
         """Decrypts an encrypted iOS backup.
         :param backup_path: Path to the encrypted backup folder
         :param dest_path: Path to the folder where to store the decrypted backup
         """
         self.backup_path = os.path.abspath(backup_path)
         self.dest_path = dest_path
+        if not 1 <= max_workers <= MAX_DECRYPT_WORKERS:
+            raise ValueError(f"max_workers must be between 1 and {MAX_DECRYPT_WORKERS}")
+        self.max_workers = max_workers
         self._backup: Optional[MVTEncryptedBackup] = None
         self._decryption_key: Optional[str] = None
 
     def can_process(self) -> bool:
         return self._backup is not None
 
+    def _process_file(
+        self,
+        *,
+        file_id: str,
+        file_bplist: bytes,
+        output_path: Path,
+        relative_path: str,
+        domain: str,
+    ) -> None:
+        assert self._backup is not None
+        self._backup.extract_file_by_id(
+            file_id=file_id,
+            file_bplist=file_bplist,
+            output_filename=str(output_path),
+        )
+        log.info(
+            "Decrypted file %s [%s] to %s/%s",
+            relative_path,
+            domain,
+            output_path.parent,
+            file_id,
+        )
+
+    @staticmethod
+    def _wait_for_files(
+        pending: dict[Future[None], str], *, all_files: bool = False
+    ) -> None:
+        if not pending:
+            return
+
+        done, _ = wait(
+            pending,
+            return_when=ALL_COMPLETED if all_files else FIRST_COMPLETED,
+        )
+        for future in done:
+            relative_path = pending.pop(future)
+            try:
+                future.result()
+            except Exception as exc:
+                log.error("Failed to decrypt file %s: %s", relative_path, exc)
+
     @staticmethod
     def is_encrypted(backup_path: str) -> bool:
         """Query Manifest.db file to see if it's encrypted or not.
@@ -226,45 +285,63 @@ class DecryptBackup:
         # modules expect.
         backup_root = Path(self.backup_path).resolve()
         dest_root = Path(self.dest_path).resolve()
-        with self._backup.manifest_db_cursor() as cur:
-            cur.execute(
-                "SELECT fileID, domain, relativePath, file FROM Files WHERE flags=1"
-            )
-            for file_id, domain, relative_path, file_bplist in cur:
-                # This may be a partial backup. Skip files from the manifest
-                # which do not exist locally.
-                source_file_path = backup_root / file_id[:2] / file_id
-                if not source_file_path.resolve().is_relative_to(backup_root):
-                    log.warning("Skipping unsafe file_id: %r", file_id)
-                    continue
-                if not os.path.exists(source_file_path):
-                    log.debug(
-                        "Skipping file %s. File not found in encrypted backup directory.",
-                        source_file_path,
-                    )
-                    continue
+        pending: dict[Future[None], str] = {}
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            with self._backup.manifest_db_cursor() as cur:
+                cur.execute(
+                    "SELECT fileID, domain, relativePath, file FROM Files WHERE flags=1"
+                )
+                for file_id, domain, relative_path, file_bplist in cur:
+                    # This may be a partial backup. Skip files from the manifest
+                    # which do not exist locally.
+                    source_file_path = backup_root / file_id[:2] / file_id
+                    if not source_file_path.resolve().is_relative_to(backup_root):
+                        log.warning("Skipping unsafe file_id: %r", file_id)
+                        continue
+                    if not os.path.exists(source_file_path):
+                        log.debug(
+                            "Skipping file %s. File not found in encrypted "
+                            "backup directory.",
+                            source_file_path,
+                        )
+                        continue
 
-                output_path = dest_root / file_id[:2] / file_id
-                if not output_path.resolve().is_relative_to(dest_root):
-                    log.warning("Skipping unsafe file_id: %r", file_id)
-                    continue
-                output_path.parent.mkdir(parents=True, exist_ok=True)
+                    output_path = dest_root / file_id[:2] / file_id
+                    if not output_path.resolve().is_relative_to(dest_root):
+                        log.warning("Skipping unsafe file_id: %r", file_id)
+                        continue
+                    output_path.parent.mkdir(parents=True, exist_ok=True)
 
-                try:
-                    self._backup.extract_file_by_id(
+                    if self.max_workers == 1:
+                        try:
+                            self._process_file(
+                                file_id=file_id,
+                                file_bplist=file_bplist,
+                                output_path=output_path,
+                                relative_path=relative_path,
+                                domain=domain,
+                            )
+                        except Exception as exc:
+                            log.error(
+                                "Failed to decrypt file %s: %s",
+                                relative_path,
+                                exc,
+                            )
+                        continue
+
+                    future = executor.submit(
+                        self._process_file,
                         file_id=file_id,
                         file_bplist=file_bplist,
-                        output_filename=str(output_path),
+                        output_path=output_path,
+                        relative_path=relative_path,
+                        domain=domain,
                     )
-                    log.info(
-                        "Decrypted file %s [%s] to %s/%s",
-                        relative_path,
-                        domain,
-                        output_path.parent,
-                        file_id,
-                    )
-                except Exception as exc:
-                    log.error("Failed to decrypt file %s: %s", relative_path, exc)
+                    pending[future] = relative_path
+                    if len(pending) >= self.max_workers:
+                        self._wait_for_files(pending)
+
+            self._wait_for_files(pending, all_files=True)
 
         # Copying over the root plist files as well.
         for file_name in os.listdir(self.backup_path):

tests/ios_backup/test_decrypt.py

@@ -3,6 +3,8 @@
 # Use of this software is governed by the MVT License 1.1 that can be found at
 #   https://license.mvt.re/1.1/
 
+import logging
+import threading
 from pathlib import Path
 
 from Crypto.Cipher import AES
@@ -114,7 +116,9 @@ def test_process_backup_rejects_unsafe_file_ids_and_destinations(mocker, tmp_pat
         Path(output_filename).write_bytes(b"decrypted")
 
     backup.extract_file_by_id.side_effect = extract_file_by_id
-    decryptor = DecryptBackup(str(backup_path), str(destination))
+    decryptor = DecryptBackup(
+        str(backup_path), str(destination), max_workers=1
+    )
     decryptor._backup = backup
 
     decryptor.process_backup()
@@ -123,3 +127,65 @@ def test_process_backup_rejects_unsafe_file_ids_and_destinations(mocker, tmp_pat
     assert not (outside / symlink_file_id).exists()
     backup.extract_file_by_id.assert_called_once()
     assert backup.extract_file_by_id.call_args.kwargs["file_id"] == safe_file_id
+
+
+def test_process_backup_decrypts_files_concurrently(mocker, tmp_path):
+    backup_path = tmp_path / "backup"
+    destination = tmp_path / "destination"
+    backup_path.mkdir()
+
+    file_ids = ["ab" + "1" * 38, "cd" + "2" * 38]
+    for file_id in file_ids:
+        source_path = backup_path / file_id[:2] / file_id
+        source_path.parent.mkdir()
+        source_path.write_bytes(b"encrypted")
+
+    cursor = mocker.MagicMock()
+    cursor.__iter__.return_value = iter(
+        (file_id, "Domain", file_id, b"plist") for file_id in file_ids
+    )
+    cursor_context = mocker.MagicMock()
+    cursor_context.__enter__.return_value = cursor
+
+    barrier = threading.Barrier(2)
+    backup = mocker.MagicMock()
+    backup.manifest_db_cursor.return_value = cursor_context
+
+    def extract_file_by_id(*, file_id, output_filename, **kwargs):
+        barrier.wait(timeout=5)
+        Path(output_filename).write_bytes(file_id.encode())
+
+    backup.extract_file_by_id.side_effect = extract_file_by_id
+    decryptor = DecryptBackup(str(backup_path), str(destination), max_workers=2)
+    decryptor._backup = backup
+
+    decryptor.process_backup()
+
+    for file_id in file_ids:
+        assert (destination / file_id[:2] / file_id).read_bytes() == file_id.encode()
+
+
+def test_process_backup_logs_worker_errors(mocker, tmp_path, caplog):
+    backup_path = tmp_path / "backup"
+    destination = tmp_path / "destination"
+    backup_path.mkdir()
+    file_id = "ef" + "3" * 38
+    source_path = backup_path / file_id[:2] / file_id
+    source_path.parent.mkdir()
+    source_path.write_bytes(b"encrypted")
+
+    cursor = mocker.MagicMock()
+    cursor.__iter__.return_value = iter([(file_id, "Domain", "failing-file", b"plist")])
+    cursor_context = mocker.MagicMock()
+    cursor_context.__enter__.return_value = cursor
+
+    backup = mocker.MagicMock()
+    backup.manifest_db_cursor.return_value = cursor_context
+    backup.extract_file_by_id.side_effect = ValueError("broken file")
+    decryptor = DecryptBackup(str(backup_path), str(destination))
+    decryptor._backup = backup
+
+    with caplog.at_level(logging.ERROR, logger="mvt.ios.decrypt"):
+        decryptor.process_backup()
+
+    assert "Failed to decrypt file failing-file: broken file" in caplog.text