From 9d71661f558c6572b6af29a3792cc15287d54d56 Mon Sep 17 00:00:00 2001 From: Janik Besendorf Date: Wed, 15 Jul 2026 09:40:18 +0200 Subject: [PATCH] Run extraction modules in parallel --- README.md | 5 + docs/development.md | 34 ++- src/mvt/android/cli.py | 21 ++ src/mvt/android/cmd_check_androidqf.py | 5 + src/mvt/android/cmd_check_backup.py | 2 + src/mvt/android/cmd_check_bugreport.py | 2 + src/mvt/android/cmd_check_intrusion_logs.py | 2 + .../android/modules/androidqf/aqf_packages.py | 9 +- src/mvt/android/modules/androidqf/base.py | 15 +- src/mvt/android/modules/backup/base.py | 33 +-- src/mvt/android/modules/bugreport/base.py | 30 ++- src/mvt/common/command.py | 224 ++++++++++++++-- src/mvt/common/help.py | 1 + src/mvt/common/log.py | 42 ++- src/mvt/common/module.py | 24 +- src/mvt/ios/cli.py | 11 + src/mvt/ios/cmd_check_backup.py | 2 + src/mvt/ios/cmd_check_fs.py | 2 + tests/common/test_command.py | 241 +++++++++++++++++- tests/test_check_ios_backup.py | 8 + tests/test_custom_modules.py | 20 +- 21 files changed, 643 insertions(+), 90 deletions(-) diff --git a/README.md b/README.md index adcee7a..9192ae6 100644 --- a/README.md +++ b/README.md @@ -82,6 +82,11 @@ Module-running `check-*` commands can load custom Python modules with [development documentation](https://docs.mvt.re/en/latest/development/) for details. +Extraction commands run independent modules concurrently with four workers by +default. Use `--jobs 1` for sequential execution or another positive value to +set the worker limit. Dependencies retain their declared ordering, and parallel +console logs are printed as grouped module blocks. + ## License diff --git a/docs/development.md b/docs/development.md index 93a612f..c3a9eda 100644 --- a/docs/development.md +++ b/docs/development.md @@ -39,6 +39,33 @@ Selecting a single module also runs its transitive dependencies. If a dependency is unavailable or the dependency graph contains a cycle, the command logs a warning and does not run any modules. +## Parallel module execution + +Extraction modules run concurrently, using four worker threads by default. Use +`--jobs INTEGER` to change the worker limit, or `--jobs 1` for sequential +execution and live per-line logging: + +```bash +mvt-ios check-backup --jobs 8 --output ./out ./backup +``` + +The scheduler starts a module only after all of its declared dependencies have +completed. Results, alerts, and timeline entries are still aggregated in stable +topological order. During parallel execution, console records from each module +are buffered and printed together as a labeled block when that module finishes; +`command.log` is written immediately and retains the complete log stream. + +Parallel-safe modules must log through `self.log`, must not write directly to +stdout or stderr, and must not mutate shared global state. A module that uses a +non-thread-safe resource or direct terminal output can opt out: + +```python +class TerminalModule(MVTModule): + parallel_safe = False +``` + +Such modules run synchronously and exclusively after active workers finish. + ## Custom modules Module-running `check-*` commands can load custom modules from Python files that @@ -82,6 +109,9 @@ class ExampleCustomModule(MVTModule): return None ``` +Custom modules are considered parallel-safe by default. Follow the parallel +module requirements above or set `parallel_safe = False`. + Use `supported_commands` to restrict a module to specific platform/command pairs. Missing or empty `supported_commands` means the module is available to all commands, which keeps older modules compatible. Supported pairs are: @@ -120,7 +150,9 @@ class DependentCustomModule(MVTModule): Some MVT modules extract and process significant amounts of data during the analysis process or while checking results against known indicators. Care must be take to avoid inefficient code paths as we add new modules. -MVT modules can be profiled with Python built-in `cProfile` by setting the `MVT_PROFILE` environment variable. +MVT modules can be profiled with Python built-in `cProfile` by setting the +`MVT_PROFILE` environment variable. Profiling forces sequential execution even +when `--jobs` is greater than one. ```bash MVT_PROFILE=1 dev/mvt-ios check-backup test_backup diff --git a/src/mvt/android/cli.py b/src/mvt/android/cli.py index 1123742..31d04ee 100644 --- a/src/mvt/android/cli.py +++ b/src/mvt/android/cli.py @@ -29,6 +29,7 @@ from mvt.common.help import ( HELP_MSG_DISABLE_UPDATE_CHECK, HELP_MSG_HASHES, HELP_MSG_IOC, + HELP_MSG_JOBS, HELP_MSG_LIST_MODULES, HELP_MSG_LOAD_MODULE, HELP_MSG_MODULE, @@ -180,6 +181,9 @@ def check_adb(ctx): help=HELP_MSG_LOAD_MODULE, ) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("BUGREPORT_PATH", type=click.Path(exists=True)) @click.pass_context def check_bugreport( @@ -190,6 +194,7 @@ def check_bugreport( module, load_module, verbose, + jobs, bugreport_path, ): set_verbose_logging(verbose) @@ -204,6 +209,7 @@ def check_bugreport( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: @@ -245,6 +251,9 @@ def check_bugreport( @click.option("--non-interactive", "-n", is_flag=True, help=HELP_MSG_NONINTERACTIVE) @click.option("--backup-password", "-p", help=HELP_MSG_ANDROID_BACKUP_PASSWORD) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("BACKUP_PATH", type=click.Path(exists=True)) @click.pass_context def check_backup( @@ -256,6 +265,7 @@ def check_backup( non_interactive, backup_password, verbose, + jobs, backup_path, ): set_verbose_logging(verbose) @@ -274,6 +284,7 @@ def check_backup( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: @@ -319,6 +330,9 @@ def check_backup( @click.option("--non-interactive", "-n", is_flag=True, help=HELP_MSG_NONINTERACTIVE) @click.option("--backup-password", "-p", help=HELP_MSG_ANDROID_BACKUP_PASSWORD) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("ANDROIDQF_PATH", type=click.Path(exists=True)) @click.pass_context def check_androidqf( @@ -334,6 +348,7 @@ def check_androidqf( non_interactive, backup_password, verbose, + jobs, androidqf_path, ): set_verbose_logging(verbose) @@ -354,6 +369,7 @@ def check_androidqf( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: @@ -405,6 +421,9 @@ def check_androidqf( ), ) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("LOGS_PATH", type=click.Path(exists=True)) @click.pass_context def check_intrusion_logs( @@ -416,6 +435,7 @@ def check_intrusion_logs( load_module, timezone, verbose, + jobs, logs_path, ): set_verbose_logging(verbose) @@ -434,6 +454,7 @@ def check_intrusion_logs( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: diff --git a/src/mvt/android/cmd_check_androidqf.py b/src/mvt/android/cmd_check_androidqf.py index 41ad802..4b4c5be 100644 --- a/src/mvt/android/cmd_check_androidqf.py +++ b/src/mvt/android/cmd_check_androidqf.py @@ -52,6 +52,7 @@ class CmdAndroidCheckAndroidQF(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -67,6 +68,7 @@ class CmdAndroidCheckAndroidQF(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "android" @@ -215,6 +217,7 @@ class CmdAndroidCheckAndroidQF(Command): hashes=self.hashes, sub_command=True, custom_modules=self.custom_modules, + jobs=self.jobs, ) cmd.from_zip(bugreport) cmd.run() @@ -245,6 +248,7 @@ class CmdAndroidCheckAndroidQF(Command): hashes=self.hashes, sub_command=True, custom_modules=self.custom_modules, + jobs=self.jobs, ) try: cmd.from_ab(backup) @@ -325,6 +329,7 @@ class CmdAndroidCheckAndroidQF(Command): hashes=self.hashes, sub_command=True, custom_modules=self.custom_modules, + jobs=self.jobs, ) cmd.run() diff --git a/src/mvt/android/cmd_check_backup.py b/src/mvt/android/cmd_check_backup.py index b75bb34..8599c51 100644 --- a/src/mvt/android/cmd_check_backup.py +++ b/src/mvt/android/cmd_check_backup.py @@ -47,6 +47,7 @@ class CmdAndroidCheckBackup(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -62,6 +63,7 @@ class CmdAndroidCheckBackup(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "android" diff --git a/src/mvt/android/cmd_check_bugreport.py b/src/mvt/android/cmd_check_bugreport.py index 1c03c6d..c6444d0 100644 --- a/src/mvt/android/cmd_check_bugreport.py +++ b/src/mvt/android/cmd_check_bugreport.py @@ -34,6 +34,7 @@ class CmdAndroidCheckBugreport(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -49,6 +50,7 @@ class CmdAndroidCheckBugreport(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "android" diff --git a/src/mvt/android/cmd_check_intrusion_logs.py b/src/mvt/android/cmd_check_intrusion_logs.py index 95f2089..a21fe6d 100644 --- a/src/mvt/android/cmd_check_intrusion_logs.py +++ b/src/mvt/android/cmd_check_intrusion_logs.py @@ -37,6 +37,7 @@ class CmdAndroidCheckIntrusionLogs(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -52,6 +53,7 @@ class CmdAndroidCheckIntrusionLogs(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "android" diff --git a/src/mvt/android/modules/androidqf/aqf_packages.py b/src/mvt/android/modules/androidqf/aqf_packages.py index 294ad53..dcc01e0 100644 --- a/src/mvt/android/modules/androidqf/aqf_packages.py +++ b/src/mvt/android/modules/androidqf/aqf_packages.py @@ -8,8 +8,6 @@ import logging import time from typing import Optional -from rich.progress import track - from mvt.android.utils import ( BROWSER_INSTALLERS, PLAY_STORE_INSTALLERS, @@ -153,10 +151,9 @@ class AQFPackages(AndroidQFModule): if total_hashes == 0: return - progress_desc = f"Looking up {total_hashes} package files on VirusTotal..." - for index, file_hash in enumerate( - track(files_by_hash, description=progress_desc) - ): + self.log.info("Looking up %d package files on VirusTotal...", total_hashes) + for index, file_hash in enumerate(files_by_hash): + self.log.debug("VirusTotal lookup %d/%d", index + 1, total_hashes) try: results = virustotal_lookup(file_hash) except VTNoKey as exc: diff --git a/src/mvt/android/modules/androidqf/base.py b/src/mvt/android/modules/androidqf/base.py index b0304d0..6241409 100644 --- a/src/mvt/android/modules/androidqf/base.py +++ b/src/mvt/android/modules/androidqf/base.py @@ -81,12 +81,13 @@ class AndroidQFModule(MVTModule): return None def _get_file_content(self, file_path): - if self.archive: - handle = self.archive.open(file_path) - else: - handle = open(os.path.join(self.parent_path, file_path), "rb") + with self.resource_lock: + if self.archive: + handle = self.archive.open(file_path) + else: + handle = open(os.path.join(self.parent_path, file_path), "rb") - data = handle.read() - handle.close() + data = handle.read() + handle.close() - return data + return data diff --git a/src/mvt/android/modules/backup/base.py b/src/mvt/android/modules/backup/base.py index 6383e4b..fc40877 100644 --- a/src/mvt/android/modules/backup/base.py +++ b/src/mvt/android/modules/backup/base.py @@ -55,21 +55,22 @@ class BackupModule(MVTModule): return fnmatch.filter(self.files, pattern) def _get_file_content(self, file_path: str) -> bytes: - handle = None - if self.tar: - try: - member = self.tar.getmember(file_path) - handle = self.tar.extractfile(member) - if not handle: - raise ValueError(f"Could not extract file: {file_path}") - except KeyError: - raise FileNotFoundError(f"File not found in tar: {file_path}") - elif self.backup_path: - handle = open(os.path.join(self.backup_path, file_path), "rb") - else: - raise ValueError("No backup path or tar file provided") + with self.resource_lock: + handle = None + if self.tar: + try: + member = self.tar.getmember(file_path) + handle = self.tar.extractfile(member) + if not handle: + raise ValueError(f"Could not extract file: {file_path}") + except KeyError: + raise FileNotFoundError(f"File not found in tar: {file_path}") + elif self.backup_path: + handle = open(os.path.join(self.backup_path, file_path), "rb") + else: + raise ValueError("No backup path or tar file provided") - data = handle.read() - handle.close() + data = handle.read() + handle.close() - return data + return data diff --git a/src/mvt/android/modules/bugreport/base.py b/src/mvt/android/modules/bugreport/base.py index 156e01c..ddd69e3 100644 --- a/src/mvt/android/modules/bugreport/base.py +++ b/src/mvt/android/modules/bugreport/base.py @@ -66,20 +66,23 @@ class BugReportModule(MVTModule): return [] def _get_file_content(self, file_path: str) -> bytes: - if self.zip_archive: - handle = self.zip_archive.open(file_path) - else: - if not self.extract_path: - raise ValueError("extract_path is not set") - joined = os.path.join(self.extract_path, file_path) - if not Path(joined).resolve().is_relative_to(Path(self.extract_path).resolve()): - raise ValueError("unsafe file_path") - handle = open(joined, "rb") + with self.resource_lock: + if self.zip_archive: + handle = self.zip_archive.open(file_path) + else: + if not self.extract_path: + raise ValueError("extract_path is not set") + joined = os.path.join(self.extract_path, file_path) + if not Path(joined).resolve().is_relative_to( + Path(self.extract_path).resolve() + ): + raise ValueError("unsafe file_path") + handle = open(joined, "rb") - data = handle.read() - handle.close() + data = handle.read() + handle.close() - return data + return data def _get_dumpstate_file(self) -> Optional[bytes]: main = self._get_files_by_pattern("main_entry.txt") @@ -102,7 +105,8 @@ class BugReportModule(MVTModule): def _get_file_modification_time(self, file_path: str) -> datetime.datetime: if self.zip_archive: - file_timetuple = self.zip_archive.getinfo(file_path).date_time + with self.resource_lock: + file_timetuple = self.zip_archive.getinfo(file_path).date_time return datetime.datetime(*file_timetuple) else: if not self.extract_path: diff --git a/src/mvt/common/command.py b/src/mvt/common/command.py index 8d21aa3..acfadd3 100644 --- a/src/mvt/common/command.py +++ b/src/mvt/common/command.py @@ -6,7 +6,10 @@ import json import logging import os +import queue import sys +import threading +from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime from heapq import heappop, heappush from typing import Any, Optional @@ -18,6 +21,7 @@ from rich.text import Text from .alerts import AlertLevel, AlertStore from .config import settings from .indicators import Indicators +from .log import MVTLogHandler, finish_module_log_buffer, start_module_log_buffer from .module import EncryptedBackupError, MVTModule, run_module, save_timeline from .module_loader import module_supports_command from .module_types import ModuleTimeline @@ -46,6 +50,7 @@ class Command: disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: self.name = "" self.platform = "" @@ -61,6 +66,10 @@ class Command: self.sub_command = sub_command self.disable_version_check = disable_version_check self.disable_indicator_check = disable_indicator_check + if jobs < 1: + raise ValueError("jobs must be at least 1") + self.jobs = jobs + self._resource_lock = threading.RLock() # This dictionary can contain options that will be passed down from # the Command to all modules. This can for example be used to pass @@ -346,35 +355,32 @@ class Command: return ordered - def run(self) -> None: - ordered_modules = self._ordered_modules() - if ordered_modules is None: - return - + def _run_module( + self, + module: type[MVTModule], + executed_by_type: dict[type[MVTModule], MVTModule], + capture_logs: bool, + ) -> tuple[MVTModule, list[logging.LogRecord], bool]: + """Initialize and run one module, optionally buffering console logs.""" + token = start_module_log_buffer() if capture_logs else None + records: list[logging.LogRecord] = [] + encrypted = False try: - self.init() - except NotImplementedError: - pass - - executed_by_type: dict[type[MVTModule], MVTModule] = {} - for module in ordered_modules: - - module_logger = logging.getLogger(module.__module__) - m = module( target_path=self.target_path, results_path=self.results_path, module_options=self.module_options, - log=module_logger, + log=logging.getLogger(module.__module__), ) m.dependency_modules = { dependency: executed_by_type[dependency] for dependency in module.dependencies } + m.resource_lock = self._resource_lock if self.iocs.total_ioc_count: + # IOC collections are shared read-only during module execution. m.indicators = self.iocs - m.indicators.log = m.log if self.serial: m.serial = self.serial @@ -387,22 +393,190 @@ class Command: try: run_module(m) except EncryptedBackupError: - self.log.critical( - "The backup appears to be encrypted. " - "Please decrypt it first using `mvt-ios decrypt-backup`." - ) - return + encrypted = True + finally: + if token is not None: + records = finish_module_log_buffer(token) - self.executed.append(m) - executed_by_type[module] = m - self.timeline.extend(m.timeline) - self.alertstore.extend(m.alertstore.alerts) + return m, records, encrypted + + @staticmethod + def _console_handler() -> Optional[MVTLogHandler]: + for handler in reversed(logging.getLogger("mvt").handlers): + if isinstance(handler, MVTLogHandler): + return handler + return None + + def _aggregate_modules( + self, + ordered_modules: list[type[MVTModule]], + completed: dict[type[MVTModule], MVTModule], + ) -> None: + """Aggregate results in stable topological order.""" + for module in ordered_modules: + if module not in completed: + continue + instance = completed[module] + self.executed.append(instance) + self.timeline.extend(instance.timeline) + self.alertstore.extend(instance.alertstore.alerts) + + def _run_sequential( + self, ordered_modules: list[type[MVTModule]] + ) -> tuple[dict[type[MVTModule], MVTModule], bool]: + completed: dict[type[MVTModule], MVTModule] = {} + for module in ordered_modules: + instance, _, encrypted = self._run_module(module, completed, False) + if encrypted: + return completed, True + completed[module] = instance + return completed, False + + def _run_parallel( + self, ordered_modules: list[type[MVTModule]], jobs: int + ) -> tuple[dict[type[MVTModule], MVTModule], bool]: + """Run a stable module DAG with bounded worker threads.""" + completed: dict[type[MVTModule], MVTModule] = {} + scheduled: set[type[MVTModule]] = set() + running: dict[Future, type[MVTModule]] = {} + completion_queue: queue.Queue[Future] = queue.Queue() + handler = self._console_handler() + status = None + fatal = False + + def ready_modules() -> list[type[MVTModule]]: + return [ + module + for module in ordered_modules + if module not in scheduled + and all(dependency in completed for dependency in module.dependencies) + ] + + def update_status() -> None: + if status is None: + return + names = ", ".join(module.__name__ for module in running.values()) + status.update( + f"Modules: {len(completed)}/{len(ordered_modules)} complete" + + (f"; running {names}" if names else "") + ) + + if handler and handler.console.is_terminal: + status = handler.console.status("") + status.start() + + try: + with ThreadPoolExecutor(max_workers=jobs) as executor: + while len(completed) < len(ordered_modules): + ready = ready_modules() + + # An unsafe module is a scheduler barrier: it starts only + # after active workers drain and blocks later ready work. + if ready and not ready[0].parallel_safe: + if running: + future = completion_queue.get() + module = running.pop(future) + instance, records, encrypted = future.result() + if handler: + handler.emit_module_records(module.__name__, records) + if encrypted: + fatal = True + break + completed[module] = instance + update_status() + continue + + module = ready[0] + scheduled.add(module) + if status is not None: + status.update( + f"Modules: {len(completed)}/{len(ordered_modules)} " + f"complete; running {module.__name__}" + ) + instance, records, encrypted = self._run_module( + module, completed, True + ) + if handler: + handler.emit_module_records(module.__name__, records) + if encrypted: + fatal = True + break + completed[module] = instance + update_status() + continue + + for module in ready: + if len(running) >= jobs or not module.parallel_safe: + break + scheduled.add(module) + future = executor.submit( + self._run_module, module, dict(completed), True + ) + running[future] = module + future.add_done_callback(completion_queue.put) + + update_status() + if not running: + break + + future = completion_queue.get() + module = running.pop(future) + instance, records, encrypted = future.result() + if handler: + handler.emit_module_records(module.__name__, records) + if encrypted: + fatal = True + break + completed[module] = instance + update_status() + + if fatal: + for future in running: + future.cancel() + # The executor context waits for already-active work. Its + # results are intentionally discarded after the fatal error. + finally: + if status is not None: + status.stop() + + return completed, fatal + + def run(self) -> None: + ordered_modules = self._ordered_modules() + if ordered_modules is None: + return + + try: + self.init() + except NotImplementedError: + pass + + jobs = self.jobs + if settings.PROFILE and jobs > 1: + self.log.warning( + "MVT_PROFILE is enabled; forcing sequential module execution." + ) + jobs = 1 + + if jobs == 1: + completed, fatal = self._run_sequential(ordered_modules) + else: + completed, fatal = self._run_parallel(ordered_modules, jobs) + + self._aggregate_modules(ordered_modules, completed) try: self.finish() except NotImplementedError: pass + if fatal: + self.log.critical( + "The backup appears to be encrypted. " + "Please decrypt it first using `mvt-ios decrypt-backup`." + ) + return + # We only store the timeline from the parent/main command if self.sub_command: return diff --git a/src/mvt/common/help.py b/src/mvt/common/help.py index 90e71ca..c08cad3 100644 --- a/src/mvt/common/help.py +++ b/src/mvt/common/help.py @@ -17,6 +17,7 @@ HELP_MSG_LOAD_MODULE = ( HELP_MSG_NONINTERACTIVE = "Don't ask interactive questions during processing" HELP_MSG_HASHES = "Generate hashes of all the files analyzed" HELP_MSG_VERBOSE = "Verbose mode" +HELP_MSG_JOBS = "Number of extraction modules to run concurrently" HELP_MSG_CHECK_IOCS = "Compare stored JSON results to provided indicators" HELP_MSG_STIX2 = "Download public STIX2 indicators" HELP_MSG_DISABLE_UPDATE_CHECK = "Disable MVT version update check" diff --git a/src/mvt/common/log.py b/src/mvt/common/log.py index 498b13d..ee1b9e3 100644 --- a/src/mvt/common/log.py +++ b/src/mvt/common/log.py @@ -4,6 +4,8 @@ # https://license.mvt.re/1.1/ import logging +import threading +from contextvars import ContextVar, Token from rich.console import Console from rich.logging import RichHandler from typing import Optional @@ -27,16 +29,33 @@ logging.addLevelName(HIGH_ALERT, "HIGH") logging.addLevelName(CRITICAL_ALERT, "CRITICAL") +_module_log_buffer: ContextVar[Optional[list[logging.LogRecord]]] = ContextVar( + "mvt_module_log_buffer", default=None +) + + +def start_module_log_buffer() -> Token: + """Capture console records in the current module execution context.""" + return _module_log_buffer.set([]) + + +def finish_module_log_buffer(token: Token) -> list[logging.LogRecord]: + records = _module_log_buffer.get() or [] + _module_log_buffer.reset(token) + return records + + class MVTLogHandler(RichHandler): def __init__(self, console: Optional[Console] = None, level: int = logging.DEBUG): super().__init__(console=console, level=level) + self._console_lock = threading.RLock() def __add_prefix_space(self, level: str) -> str: max_length = len("CRITICAL ALERT") space = max_length - len(level) return f"{level}{' ' * space}" - def emit(self, record: logging.LogRecord): + def _emit_record(self, record: logging.LogRecord) -> None: try: msg = rf"[grey50]\[{record.name}][/] {self.format(record)}" @@ -63,3 +82,24 @@ class MVTLogHandler(RichHandler): except Exception: self.handleError(record) + + def emit(self, record: logging.LogRecord): + buffer = _module_log_buffer.get() + if buffer is not None: + # More than one CLI can be imported in-process (notably in tests), + # which can install multiple console handlers for the same record. + if not buffer or buffer[-1] is not record: + buffer.append(record) + return + + with self._console_lock: + self._emit_record(record) + + def emit_module_records( + self, module_name: str, records: list[logging.LogRecord] + ) -> None: + """Print one module's captured records without interleaving.""" + with self._console_lock: + self.console.print(f"[bold cyan]--- {module_name} ---[/bold cyan]") + for record in records: + self._emit_record(record) diff --git a/src/mvt/common/module.py b/src/mvt/common/module.py index 14c5fe5..e597f15 100644 --- a/src/mvt/common/module.py +++ b/src/mvt/common/module.py @@ -8,6 +8,7 @@ import json import logging import os import re +import threading from dataclasses import asdict, is_dataclass from typing import Any, Dict, Optional, Sequence @@ -45,6 +46,7 @@ class MVTModule: slug: Optional[str] = None dependencies: Sequence[type["MVTModule"]] = () supported_commands: Sequence[tuple[str, str]] = () + parallel_safe: bool = True def __init__( self, @@ -83,6 +85,9 @@ class MVTModule: self.results: ModuleResults = results if results is not None else [] self.timeline: ModuleTimeline = [] self.dependency_modules: Dict[type["MVTModule"], "MVTModule"] = {} + # Commands replace this with a command-scoped lock for modules sharing + # archive handles or other non-thread-safe read resources. + self.resource_lock = threading.RLock() def get_dependency_results( self, module_class: type["MVTModule"] @@ -161,17 +166,20 @@ class MVTModule: """ timeline_set = set() + deduplicated = [] for record in timeline: - timeline_set.add( - json.dumps( - asdict(record) - if is_dataclass(record) and not isinstance(record, type) - else record, - sort_keys=True, - ) + serialized = json.dumps( + asdict(record) + if is_dataclass(record) and not isinstance(record, type) + else record, + sort_keys=True, ) + if serialized in timeline_set: + continue + timeline_set.add(serialized) + deduplicated.append(json.loads(serialized)) - return [json.loads(record) for record in timeline_set] + return deduplicated def to_timeline(self) -> None: """Convert results into a timeline.""" diff --git a/src/mvt/ios/cli.py b/src/mvt/ios/cli.py index 3f4beef..db5404c 100644 --- a/src/mvt/ios/cli.py +++ b/src/mvt/ios/cli.py @@ -38,6 +38,7 @@ from mvt.common.help import ( HELP_MSG_LOAD_MODULE, HELP_MSG_MODULE, HELP_MSG_VERBOSE, + HELP_MSG_JOBS, HELP_MSG_CHECK_FS, HELP_MSG_CHECK_IOCS, HELP_MSG_STIX2, @@ -288,6 +289,9 @@ def extract_key(password, key_file, backup_path): ) @click.option("--hashes", "-H", is_flag=True, help=HELP_MSG_HASHES) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("BACKUP_PATH", type=click.Path(exists=True)) @click.pass_context def check_backup( @@ -300,6 +304,7 @@ def check_backup( load_module, hashes, verbose, + jobs, backup_path, ): set_verbose_logging(verbose) @@ -316,6 +321,7 @@ def check_backup( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: @@ -357,6 +363,9 @@ def check_backup( ) @click.option("--hashes", "-H", is_flag=True, help=HELP_MSG_HASHES) @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE) +@click.option( + "--jobs", type=click.IntRange(min=1), default=4, show_default=True, help=HELP_MSG_JOBS +) @click.argument("DUMP_PATH", type=click.Path(exists=True)) @click.pass_context def check_fs( @@ -369,6 +378,7 @@ def check_fs( load_module, hashes, verbose, + jobs, dump_path, ): set_verbose_logging(verbose) @@ -385,6 +395,7 @@ def check_fs( disable_version_check=_get_disable_flags(ctx)[0], disable_indicator_check=_get_disable_flags(ctx)[1], custom_modules=custom_modules, + jobs=jobs, ) if list_modules: diff --git a/src/mvt/ios/cmd_check_backup.py b/src/mvt/ios/cmd_check_backup.py index 1b90584..1ebae32 100644 --- a/src/mvt/ios/cmd_check_backup.py +++ b/src/mvt/ios/cmd_check_backup.py @@ -38,6 +38,7 @@ class CmdIOSCheckBackup(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -53,6 +54,7 @@ class CmdIOSCheckBackup(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "ios" diff --git a/src/mvt/ios/cmd_check_fs.py b/src/mvt/ios/cmd_check_fs.py index e76146e..860ff48 100644 --- a/src/mvt/ios/cmd_check_fs.py +++ b/src/mvt/ios/cmd_check_fs.py @@ -31,6 +31,7 @@ class CmdIOSCheckFS(Command): disable_version_check: bool = False, disable_indicator_check: bool = False, custom_modules: Optional[list[type[MVTModule]]] = None, + jobs: int = 4, ) -> None: super().__init__( target_path=target_path, @@ -45,6 +46,7 @@ class CmdIOSCheckFS(Command): disable_version_check=disable_version_check, disable_indicator_check=disable_indicator_check, custom_modules=custom_modules, + jobs=jobs, ) self.platform = "ios" diff --git a/tests/common/test_command.py b/tests/common/test_command.py index 865f9f1..858a3b4 100644 --- a/tests/common/test_command.py +++ b/tests/common/test_command.py @@ -5,9 +5,16 @@ import json import logging +import threading +import time +from io import StringIO + +from rich.console import Console from mvt.common.command import Command -from mvt.common.module import MVTModule +from mvt.common.config import settings +from mvt.common.log import MVTLogHandler +from mvt.common.module import EncryptedBackupError, MVTModule class RecordingModule(MVTModule): @@ -83,8 +90,17 @@ class TestCommand: alerts = json.loads((tmp_path / "alerts.json").read_text()) assert alerts[0]["event"]["payload"] == "\\xa8\\xa9" + def test_timeline_deduplication_preserves_first_seen_order(self): + timeline = [ + {"timestamp": "same", "event": "first"}, + {"timestamp": "same", "event": "second"}, + {"timestamp": "same", "event": "first"}, + ] + + assert MVTModule._deduplicate_timeline(timeline) == timeline[:2] + def test_modules_run_in_stable_topological_order(self): - cmd = RecordingCommand() + cmd = RecordingCommand(jobs=1) cmd.modules = [ThirdModule, IndependentModule, SecondModule, FirstModule] cmd.run() @@ -99,7 +115,7 @@ class TestCommand: assert second.results == ["first", "second"] def test_selected_module_runs_transitive_dependencies(self): - cmd = RecordingCommand(module_name="ThirdModule") + cmd = RecordingCommand(module_name="ThirdModule", jobs=1) cmd.modules = [ThirdModule, SecondModule, FirstModule, IndependentModule] cmd.run() @@ -159,7 +175,7 @@ class TestCommand: ] def test_selected_custom_module_runs(self): - cmd = RecordingCommand(module_name="CustomIOSBackupModule") + cmd = RecordingCommand(module_name="CustomIOSBackupModule", jobs=1) cmd.platform = "ios" cmd.name = "check-backup" cmd.custom_modules = [CustomIOSBackupModule] @@ -179,7 +195,7 @@ class TestCommand: assert RecordingModule.run_order == [] def test_custom_module_dependencies_use_topological_order(self): - cmd = RecordingCommand(module_name="CustomDependsOnBuiltin") + cmd = RecordingCommand(module_name="CustomDependsOnBuiltin", jobs=1) cmd.platform = "ios" cmd.name = "check-backup" cmd.modules = [SecondModule, FirstModule] @@ -188,3 +204,218 @@ class TestCommand: cmd.run() assert RecordingModule.run_order == ["FirstModule", "CustomDependsOnBuiltin"] + + def test_parallel_modules_overlap_and_respect_worker_limit(self): + lock = threading.Lock() + barrier = threading.Barrier(2) + active = 0 + maximum = 0 + overlapped = [] + + class ParallelModule(RecordingModule): + def run(self): + nonlocal active, maximum + with lock: + active += 1 + maximum = max(maximum, active) + try: + barrier.wait(timeout=2) + overlapped.append(self.__class__.__name__) + time.sleep(0.02) + finally: + with lock: + active -= 1 + + class ParallelOne(ParallelModule): + pass + + class ParallelTwo(ParallelModule): + pass + + class ParallelThree(ParallelModule): + pass + + cmd = RecordingCommand(jobs=2) + cmd.modules = [ParallelOne, ParallelTwo, ParallelThree] + cmd.run() + + assert maximum == 2 + assert {"ParallelOne", "ParallelTwo"}.issubset(overlapped) + + def test_dependencies_wait_and_transitive_dependencies_run_once(self): + prerequisite_finished = threading.Event() + runs = [] + + class Prerequisite(RecordingModule): + def run(self): + runs.append("prerequisite") + self.results = ["ready"] + prerequisite_finished.set() + + class Dependent(RecordingModule): + dependencies = (Prerequisite,) + + def run(self): + assert prerequisite_finished.is_set() + assert self.get_dependency_results(Prerequisite) == ["ready"] + runs.append(self.__class__.__name__) + + class DependentOne(Dependent): + pass + + class DependentTwo(Dependent): + pass + + cmd = RecordingCommand(jobs=3) + cmd.modules = [DependentOne, DependentTwo, Prerequisite] + cmd.run() + + assert runs.count("prerequisite") == 1 + assert set(runs[1:]) == {"DependentOne", "DependentTwo"} + + def test_parallel_unsafe_module_runs_exclusively(self): + lock = threading.Lock() + active = 0 + unsafe_ran_exclusively = [] + + class SafeOne(RecordingModule): + def run(self): + nonlocal active + with lock: + active += 1 + time.sleep(0.03) + with lock: + active -= 1 + + class Unsafe(RecordingModule): + parallel_safe = False + + def run(self): + nonlocal active + with lock: + unsafe_ran_exclusively.append(active == 0) + active += 1 + time.sleep(0.01) + with lock: + active -= 1 + + class SafeTwo(SafeOne): + pass + + cmd = RecordingCommand(jobs=3) + cmd.modules = [SafeOne, Unsafe, SafeTwo] + cmd.run() + + assert unsafe_ran_exclusively == [True] + + def test_results_are_aggregated_in_topological_order(self): + release_first = threading.Event() + + class SlowFirst(RecordingModule): + def run(self): + release_first.wait(timeout=2) + self.results = ["first"] + self.timeline = [{"module": "first"}] + self.alertstore.info("first", "", {"order": 1}) + + class FastSecond(RecordingModule): + def run(self): + self.results = ["second"] + self.timeline = [{"module": "second"}] + self.alertstore.info("second", "", {"order": 2}) + release_first.set() + + cmd = RecordingCommand(jobs=2) + cmd.modules = [SlowFirst, FastSecond] + cmd.run() + + assert [type(module) for module in cmd.executed] == [SlowFirst, FastSecond] + assert [entry["module"] for entry in cmd.timeline] == ["first", "second"] + assert [alert.message for alert in cmd.alertstore.alerts] == ["first", "second"] + + def test_parallel_console_logs_are_grouped_and_file_log_is_complete( + self, tmp_path + ): + output = StringIO() + handler = MVTLogHandler( + console=Console(file=output, force_terminal=False, color_system=None) + ) + handler.setFormatter(logging.Formatter("%(message)s")) + logger = logging.getLogger("mvt") + logger.setLevel(logging.DEBUG) + logger.addHandler(handler) + release_slow = threading.Event() + + class SlowLog(RecordingModule): + def run(self): + self.log.info("slow-a") + release_slow.wait(timeout=2) + self.log.info("slow-b") + + class FastLog(RecordingModule): + def run(self): + self.log.info("fast-a") + self.log.info("fast-b") + release_slow.set() + + SlowLog.__module__ = "mvt.tests.parallel" + FastLog.__module__ = "mvt.tests.parallel" + try: + cmd = RecordingCommand(results_path=str(tmp_path), jobs=2) + cmd.modules = [SlowLog, FastLog] + cmd.run() + finally: + logger.removeHandler(handler) + + console_output = output.getvalue() + assert console_output.index("--- FastLog ---") < console_output.index( + "--- SlowLog ---" + ) + assert console_output.index("fast-a") < console_output.index("fast-b") + assert "slow-a" in console_output and "slow-b" in console_output + file_output = (tmp_path / "command.log").read_text() + for message in ("slow-a", "slow-b", "fast-a", "fast-b"): + assert message in file_output + + def test_encrypted_backup_stops_scheduling_and_cleans_up(self, caplog): + finish_called = [] + + class Encrypted(RecordingModule): + def run(self): + raise EncryptedBackupError + + class NeverScheduled(RecordingModule): + dependencies = (Encrypted,) + + class CleanupCommand(RecordingCommand): + def finish(self): + finish_called.append(True) + + cmd = CleanupCommand(jobs=2) + cmd.modules = [Encrypted, NeverScheduled] + with caplog.at_level(logging.CRITICAL): + cmd.run() + + assert finish_called == [True] + assert not any(isinstance(module, NeverScheduled) for module in cmd.executed) + assert "backup appears to be encrypted" in caplog.text + + def test_profiling_forces_sequential_execution(self, monkeypatch, caplog): + run_order = [] + + class ProfileOne(RecordingModule): + def run(self): + run_order.append("one") + + class ProfileTwo(RecordingModule): + def run(self): + run_order.append("two") + + monkeypatch.setattr(settings, "PROFILE", True) + cmd = RecordingCommand(jobs=4) + cmd.modules = [ProfileOne, ProfileTwo] + with caplog.at_level(logging.WARNING): + cmd.run() + + assert run_order == ["one", "two"] + assert "forcing sequential module execution" in caplog.text diff --git a/tests/test_check_ios_backup.py b/tests/test_check_ios_backup.py index eb10b13..fa799d5 100644 --- a/tests/test_check_ios_backup.py +++ b/tests/test_check_ios_backup.py @@ -33,3 +33,11 @@ class TestCheckBackupCommand: result = runner.invoke(check_backup, [str(tmp_path)]) assert result.exit_code == 1 assert "does not appear to be an iTunes backup folder" in caplog.text + + def test_check_rejects_non_positive_jobs(self): + result = CliRunner().invoke( + check_backup, ["--jobs", "0", get_ios_backup_folder()] + ) + + assert result.exit_code == 2 + assert "Invalid value for '--jobs'" in result.output diff --git a/tests/test_custom_modules.py b/tests/test_custom_modules.py index a00faea..fd3bfc8 100644 --- a/tests/test_custom_modules.py +++ b/tests/test_custom_modules.py @@ -138,15 +138,19 @@ def test_androidqf_propagates_custom_modules_to_nested_commands(tmp_path, monkey cmd = CmdAndroidCheckAndroidQF( target_path=str(tmp_path), custom_modules=custom_modules, + jobs=7, ) def record_available(name): def _record(command): - records[name] = [ - module.__name__ - for module in command._available_modules() - if module.__name__.startswith("Nested") - ] + records[name] = { + "jobs": command.jobs, + "modules": [ + module.__name__ + for module in command._available_modules() + if module.__name__.startswith("Nested") + ], + } return _record @@ -189,7 +193,7 @@ def test_androidqf_propagates_custom_modules_to_nested_commands(tmp_path, monkey assert cmd.run_backup_cmd() assert cmd.run_intrusion_logs_cmd() assert records == { - "bugreport": ["NestedBugreportModule"], - "backup": ["NestedBackupModule"], - "intrusion_logs": ["NestedIntrusionLogsModule"], + "bugreport": {"jobs": 7, "modules": ["NestedBugreportModule"]}, + "backup": {"jobs": 7, "modules": ["NestedBackupModule"]}, + "intrusion_logs": {"jobs": 7, "modules": ["NestedIntrusionLogsModule"]}, }