diff --git a/README.md b/README.md index adcee7ac5..9192ae68b 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 93a612f2a..c3a9eda54 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 1123742a1..31d04ee91 100644 --- a/src/mvt/android/cli.py +++ b/src/mvt/android/cli.py @@ -29,6 +29,7 @@ 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 41ad802a2..4b4c5be78 100644 --- a/src/mvt/android/cmd_check_androidqf.py +++ b/src/mvt/android/cmd_check_androidqf.py @@ -52,6 +52,7 @@ def __init__( 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 @@ def __init__( 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 @@ def run_bugreport_cmd(self) -> bool: hashes=self.hashes, sub_command=True, custom_modules=self.custom_modules, + jobs=self.jobs, ) cmd.from_zip(bugreport) cmd.run() @@ -245,6 +248,7 @@ def run_backup_cmd(self) -> bool: hashes=self.hashes, sub_command=True, custom_modules=self.custom_modules, + jobs=self.jobs, ) try: cmd.from_ab(backup) @@ -325,6 +329,7 @@ def run_intrusion_logs_cmd(self) -> bool: 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 b75bb3464..8599c51db 100644 --- a/src/mvt/android/cmd_check_backup.py +++ b/src/mvt/android/cmd_check_backup.py @@ -47,6 +47,7 @@ def __init__( 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 @@ def __init__( 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 1c03c6d2c..c6444d042 100644 --- a/src/mvt/android/cmd_check_bugreport.py +++ b/src/mvt/android/cmd_check_bugreport.py @@ -34,6 +34,7 @@ def __init__( 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 @@ def __init__( 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 95f20896d..a21fe6d7d 100644 --- a/src/mvt/android/cmd_check_intrusion_logs.py +++ b/src/mvt/android/cmd_check_intrusion_logs.py @@ -37,6 +37,7 @@ def __init__( 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 @@ def __init__( 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 294ad5302..dcc01e09b 100644 --- a/src/mvt/android/modules/androidqf/aqf_packages.py +++ b/src/mvt/android/modules/androidqf/aqf_packages.py @@ -8,8 +8,6 @@ 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 @@ def check_virustotal(self, delay: int = 0) -> None: 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 b0304d0ce..6241409dc 100644 --- a/src/mvt/android/modules/androidqf/base.py +++ b/src/mvt/android/modules/androidqf/base.py @@ -81,12 +81,13 @@ def _get_device_timezone(self): 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 6383e4bcf..fc408779c 100644 --- a/src/mvt/android/modules/backup/base.py +++ b/src/mvt/android/modules/backup/base.py @@ -55,21 +55,22 @@ def _get_files_by_pattern(self, pattern: str) -> list: 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 156e01cf2..2d49fe562 100644 --- a/src/mvt/android/modules/bugreport/base.py +++ b/src/mvt/android/modules/bugreport/base.py @@ -4,15 +4,36 @@ # https://github.com/mvt-project/mvt/blob/main/LICENSE import datetime import fnmatch +import io import logging import os +import re from pathlib import Path from typing import List, Optional from zipfile import ZipFile +from mvt.android.artifacts.artifact import AndroidArtifact from mvt.common.module import ModuleResults, MVTModule +_DUMPSTATE_BYTES_CACHE_KEY = ("android_bugreport", "dumpstate", "bytes") +_DUMPSTATE_SECTIONS_CACHE_KEY = ("android_bugreport", "dumpstate", "sections") +_SYSTEM_PROPERTIES_SECTION = b"------ SYSTEM PROPERTIES" +_DUMPSYS_SEPARATORS = { + b"DUMP OF SERVICE accessibility:", + b"DUMP OF SERVICE adb:", + b"DUMP OF SERVICE appops:", + b"DUMP OF SERVICE batterystats:", + b"DUMP OF SERVICE dbinfo:", + b"DUMP OF SERVICE package:", + b"DUMP OF SERVICE platform_compat:", +} +_DUMPSTATE_DELIMITER = ( + b"------------------------------------------------------------------------------" +) +_GETPROP_LINE_RE = re.compile(rb"\[(.+?)\]: \[(.+?)\]") + + class BugReportModule(MVTModule): """This class provides a base for all Android Bug Report modules.""" @@ -66,43 +87,160 @@ def _get_files_by_patterns(self, patterns: list) -> list: 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") - - data = handle.read() - handle.close() - - return data + 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() + + return data def _get_dumpstate_file(self) -> Optional[bytes]: - main = self._get_files_by_pattern("main_entry.txt") - if main: - main_content = self._get_file_content(main[0]) - try: - return self._get_file_content(main_content.decode().strip()) - except KeyError: - return None - - dumpstate_logs = self._get_files_by_pattern("dumpState_*.log") - if dumpstate_logs: - return self._get_file_content(dumpstate_logs[0]) - - dumpsys_files = self._get_files_by_pattern("*/dumpsys.txt") - if dumpsys_files: - return self._get_file_content(dumpsys_files[0]) - - return None + with self.resource_lock: + if _DUMPSTATE_BYTES_CACHE_KEY in self.resource_cache: + return self.resource_cache[_DUMPSTATE_BYTES_CACHE_KEY] + + content = None + main = self._get_files_by_pattern("main_entry.txt") + if main: + main_content = self._get_file_content(main[0]) + try: + content = self._get_file_content(main_content.decode().strip()) + except KeyError: + pass + else: + dumpstate_logs = self._get_files_by_pattern("dumpState_*.log") + if dumpstate_logs: + content = self._get_file_content(dumpstate_logs[0]) + else: + dumpsys_files = self._get_files_by_pattern("*/dumpsys.txt") + if dumpsys_files: + content = self._get_file_content(dumpsys_files[0]) + + self.resource_cache[_DUMPSTATE_BYTES_CACHE_KEY] = content + return content + + def _get_dumpstate_sections(self) -> dict[bytes, bytes]: + with self.resource_lock: + if _DUMPSTATE_SECTIONS_CACHE_KEY in self.resource_cache: + return self.resource_cache[_DUMPSTATE_SECTIONS_CACHE_KEY] + + section_lines: dict[bytes, list[bytes]] = { + separator: [] for separator in _DUMPSYS_SEPARATORS + } + section_lines[_SYSTEM_PROPERTIES_SECTION] = [] + current_section = None + completed_sections = set() + in_system_properties = False + content = self._get_dumpstate_file() + + if content: + for raw_line in io.BytesIO(content): + line = raw_line.rstrip(b"\r\n") + stripped = line.strip() + + if stripped.startswith(_SYSTEM_PROPERTIES_SECTION): + in_system_properties = True + elif in_system_properties: + if stripped == b"------": + in_system_properties = False + elif _GETPROP_LINE_RE.search(line): + section_lines[_SYSTEM_PROPERTIES_SECTION].append(line) + + if stripped in _DUMPSYS_SEPARATORS: + if current_section is not None: + completed_sections.add(current_section) + current_section = ( + None if stripped in completed_sections else stripped + ) + continue + + if stripped.startswith(_DUMPSTATE_DELIMITER): + if current_section is not None: + completed_sections.add(current_section) + current_section = None + continue + + if current_section is not None: + section_lines[current_section].append(line) + + sections = { + separator: b"\n".join(lines) + for separator, lines in section_lines.items() + } + self.resource_cache[_DUMPSTATE_SECTIONS_CACHE_KEY] = sections + return sections + + def _get_dumpsys_section_bytes(self, separator: bytes) -> Optional[bytes]: + sections = self._get_dumpstate_sections() + if separator in sections: + return sections[separator] + + content = self._get_dumpstate_file() + if not content: + return None + return AndroidArtifact.extract_dumpsys_section(content, separator, binary=True) + + def _get_dumpsys_section( + self, separator: str, errors: str = "replace" + ) -> Optional[str]: + cache_key = ( + "android_bugreport", + "dumpstate", + "section", + separator, + errors, + ) + with self.resource_lock: + if cache_key in self.resource_cache: + return self.resource_cache[cache_key] + + section_bytes = self._get_dumpsys_section_bytes(separator.encode("utf-8")) + section = ( + section_bytes.decode("utf-8", errors=errors) + if section_bytes is not None + else None + ) + self.resource_cache[cache_key] = section + return section + + def _get_system_properties_text(self, errors: str = "ignore") -> Optional[str]: + cache_key = ( + "android_bugreport", + "dumpstate", + "system_properties", + errors, + ) + with self.resource_lock: + if cache_key in self.resource_cache: + return self.resource_cache[cache_key] + + sections = self._get_dumpstate_sections() + section_bytes = sections.get(_SYSTEM_PROPERTIES_SECTION) + section = ( + section_bytes.decode("utf-8", errors=errors) + if section_bytes is not None + else None + ) + self.resource_cache[cache_key] = section + return section 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/android/modules/bugreport/dumpsys_accessibility.py b/src/mvt/android/modules/bugreport/dumpsys_accessibility.py index 72208c02e..09df4e3db 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_accessibility.py +++ b/src/mvt/android/modules/bugreport/dumpsys_accessibility.py @@ -34,18 +34,16 @@ def __init__( ) def run(self) -> None: - full_dumpsys = self._get_dumpstate_file() - if not full_dumpsys: + content = self._get_dumpsys_section( + "DUMP OF SERVICE accessibility:", errors="ignore" + ) + if content is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - content = self.extract_dumpsys_section( - full_dumpsys.decode("utf-8", errors="ignore"), - "DUMP OF SERVICE accessibility:", - ) self.parse(content) for result in self.results: diff --git a/src/mvt/android/modules/bugreport/dumpsys_activities.py b/src/mvt/android/modules/bugreport/dumpsys_activities.py index 2800e95ea..b5959af1f 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_activities.py +++ b/src/mvt/android/modules/bugreport/dumpsys_activities.py @@ -38,19 +38,14 @@ def __init__( self.results = results if results else [] def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + section = self._get_dumpsys_section("DUMP OF SERVICE package:", errors="ignore") + if section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - # Extract package section - section = self.extract_dumpsys_section( - content.decode("utf-8", errors="ignore"), "DUMP OF SERVICE package:" - ) - # Parse self.parse(section) diff --git a/src/mvt/android/modules/bugreport/dumpsys_adb_state.py b/src/mvt/android/modules/bugreport/dumpsys_adb_state.py index 506af30d8..7b9cb5cdd 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_adb_state.py +++ b/src/mvt/android/modules/bugreport/dumpsys_adb_state.py @@ -34,19 +34,14 @@ def __init__( ) def run(self) -> None: - full_dumpsys = self._get_dumpstate_file() - if not full_dumpsys: + content = self._get_dumpsys_section_bytes(b"DUMP OF SERVICE adb:") + if content is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - content = self.extract_dumpsys_section( - full_dumpsys, - b"DUMP OF SERVICE adb:", - binary=True, - ) self.parse(content) if self.results: self.log.info( diff --git a/src/mvt/android/modules/bugreport/dumpsys_appops.py b/src/mvt/android/modules/bugreport/dumpsys_appops.py index 91122cb57..8edef1278 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_appops.py +++ b/src/mvt/android/modules/bugreport/dumpsys_appops.py @@ -34,17 +34,14 @@ def __init__( ) def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + section = self._get_dumpsys_section("DUMP OF SERVICE appops:") + if section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - section = self.extract_dumpsys_section( - content.decode("utf-8", errors="replace"), "DUMP OF SERVICE appops:" - ) self.parse(section) self.log.info( diff --git a/src/mvt/android/modules/bugreport/dumpsys_battery_daily.py b/src/mvt/android/modules/bugreport/dumpsys_battery_daily.py index a7c0c722d..5a55bec4e 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_battery_daily.py +++ b/src/mvt/android/modules/bugreport/dumpsys_battery_daily.py @@ -34,17 +34,14 @@ def __init__( ) def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + dumpsys_section = self._get_dumpsys_section("DUMP OF SERVICE batterystats:") + if dumpsys_section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - dumpsys_section = self.extract_dumpsys_section( - content.decode("utf-8", errors="replace"), "DUMP OF SERVICE batterystats:" - ) self.parse(dumpsys_section) self.log.info("Extracted a total of %d battery daily stats", len(self.results)) diff --git a/src/mvt/android/modules/bugreport/dumpsys_battery_history.py b/src/mvt/android/modules/bugreport/dumpsys_battery_history.py index 42d395df2..236e9b7c1 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_battery_history.py +++ b/src/mvt/android/modules/bugreport/dumpsys_battery_history.py @@ -34,17 +34,14 @@ def __init__( ) def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + dumpsys_section = self._get_dumpsys_section("DUMP OF SERVICE batterystats:") + if dumpsys_section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - dumpsys_section = self.extract_dumpsys_section( - content.decode("utf-8", errors="replace"), "DUMP OF SERVICE batterystats:" - ) self.parse(dumpsys_section) self.log.info( diff --git a/src/mvt/android/modules/bugreport/dumpsys_dbinfo.py b/src/mvt/android/modules/bugreport/dumpsys_dbinfo.py index 13ba8b33b..c8f573aad 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_dbinfo.py +++ b/src/mvt/android/modules/bugreport/dumpsys_dbinfo.py @@ -36,18 +36,14 @@ def __init__( ) def run(self) -> None: - data = self._get_dumpstate_file() - if not data: + section = self._get_dumpsys_section("DUMP OF SERVICE dbinfo:", errors="ignore") + if section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - section = self.extract_dumpsys_section( - data.decode("utf-8", errors="ignore"), "DUMP OF SERVICE dbinfo:" - ) - self.parse(section) self.log.info( "Extracted a total of %d database connection pool records", diff --git a/src/mvt/android/modules/bugreport/dumpsys_getprop.py b/src/mvt/android/modules/bugreport/dumpsys_getprop.py index 198a1d2ab..1b54fd099 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_getprop.py +++ b/src/mvt/android/modules/bugreport/dumpsys_getprop.py @@ -36,29 +36,13 @@ def __init__( self.results = [] if not results else results def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + content = self._get_system_properties_text() + if content is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - lines = [] - in_getprop = False - - for line in content.decode(errors="ignore").splitlines(): - if line.strip().startswith("------ SYSTEM PROPERTIES"): - in_getprop = True - continue - - if not in_getprop: - continue - - if line.strip() == "------": - break - - lines.append(line) - - self.parse("\n".join(lines)) + self.parse(content) self.log.info("Extracted %d Android system properties", len(self.results)) diff --git a/src/mvt/android/modules/bugreport/dumpsys_packages.py b/src/mvt/android/modules/bugreport/dumpsys_packages.py index aebec4c26..186327fdb 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_packages.py +++ b/src/mvt/android/modules/bugreport/dumpsys_packages.py @@ -35,17 +35,14 @@ def __init__( ) def run(self) -> None: - data = self._get_dumpstate_file() - if not data: + content = self._get_dumpsys_section("DUMP OF SERVICE package:") + if content is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - content = self.extract_dumpsys_section( - data.decode("utf-8", errors="replace"), "DUMP OF SERVICE package:" - ) self.parse(content) for result in self.results: diff --git a/src/mvt/android/modules/bugreport/dumpsys_platform_compat.py b/src/mvt/android/modules/bugreport/dumpsys_platform_compat.py index 968bc25ea..030e85989 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_platform_compat.py +++ b/src/mvt/android/modules/bugreport/dumpsys_platform_compat.py @@ -34,18 +34,14 @@ def __init__( ) def run(self) -> None: - data = self._get_dumpstate_file() - if not data: + content = self._get_dumpsys_section("DUMP OF SERVICE platform_compat:") + if content is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - decoded_data = data.decode("utf-8", errors="replace") - content = self.extract_dumpsys_section( - decoded_data, "DUMP OF SERVICE platform_compat:" - ) self.parse(content) self.log.info("Found %d uninstalled apps", len(self.results)) diff --git a/src/mvt/android/modules/bugreport/dumpsys_receivers.py b/src/mvt/android/modules/bugreport/dumpsys_receivers.py index 8907d5e36..ef53f5fc8 100644 --- a/src/mvt/android/modules/bugreport/dumpsys_receivers.py +++ b/src/mvt/android/modules/bugreport/dumpsys_receivers.py @@ -36,17 +36,13 @@ def __init__( self.results = results if results else {} def run(self) -> None: - content = self._get_dumpstate_file() - if not content: + dumpsys_section = self._get_dumpsys_section("DUMP OF SERVICE package:") + if dumpsys_section is None: self.log.error( "Unable to find dumpstate file. " "Did you provide a valid bug report archive?" ) return - dumpsys_section = self.extract_dumpsys_section( - content.decode("utf-8", errors="replace"), "DUMP OF SERVICE package:" - ) - self.parse(dumpsys_section) self.log.info("Extracted receivers for %d intents", len(self.results)) diff --git a/src/mvt/common/command.py b/src/mvt/common/command.py index 8d21aa326..6938271c5 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 .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 @@ def __init__( 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,11 @@ def __init__( 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() + self._resource_cache: dict[object, Any] = {} # 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 +356,33 @@ def _ordered_modules(self) -> Optional[list[type[MVTModule]]]: 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 + m.resource_cache = self._resource_cache 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 +395,190 @@ def run(self) -> None: try: run_module(m) except EncryptedBackupError: - self.log.critical( - "The backup appears to be encrypted. " - "Please decrypt it first using `mvt-ios decrypt-backup`." - ) + encrypted = True + finally: + if token is not None: + records = finish_module_log_buffer(token) + + 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 - self.executed.append(m) - executed_by_type[module] = m - self.timeline.extend(m.timeline) - self.alertstore.extend(m.alertstore.alerts) + 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 90e71ca90..c08cad355 100644 --- a/src/mvt/common/help.py +++ b/src/mvt/common/help.py @@ -17,6 +17,7 @@ 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 498b13d33..ee1b9e32d 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(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 @@ def emit(self, record: logging.LogRecord): 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 14c5fe582..e7a7e1c70 100644 --- a/src/mvt/common/module.py +++ b/src/mvt/common/module.py @@ -8,6 +8,7 @@ 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,12 @@ def __init__( 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() + # Commands also replace this with a command-scoped cache. Cached values + # must be immutable, because parallel modules can read them concurrently. + self.resource_cache: Dict[object, Any] = {} def get_dependency_results( self, module_class: type["MVTModule"] @@ -161,17 +169,20 @@ def _deduplicate_timeline(timeline: list) -> list: """ 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 3f4beef3f..db5404c26 100644 --- a/src/mvt/ios/cli.py +++ b/src/mvt/ios/cli.py @@ -38,6 +38,7 @@ 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 1b90584b6..1ebae3209 100644 --- a/src/mvt/ios/cmd_check_backup.py +++ b/src/mvt/ios/cmd_check_backup.py @@ -38,6 +38,7 @@ def __init__( 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 @@ def __init__( 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 e76146ec2..860ff485d 100644 --- a/src/mvt/ios/cmd_check_fs.py +++ b/src/mvt/ios/cmd_check_fs.py @@ -31,6 +31,7 @@ def __init__( 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 @@ def __init__( 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/android_bugreport/test_bugreport.py b/tests/android_bugreport/test_bugreport.py index 2bfaef54e..220d3e9d8 100644 --- a/tests/android_bugreport/test_bugreport.py +++ b/tests/android_bugreport/test_bugreport.py @@ -4,8 +4,13 @@ # https://license.mvt.re/1.1/ import os +import threading +from concurrent.futures import ThreadPoolExecutor +from io import BytesIO from pathlib import Path +from unittest.mock import MagicMock +from mvt.android.modules.bugreport.base import BugReportModule from mvt.android.modules.bugreport.dumpsys_appops import DumpsysAppops from mvt.android.modules.bugreport.dumpsys_getprop import DumpsysGetProp from mvt.android.modules.bugreport.dumpsys_packages import DumpsysPackages @@ -90,3 +95,53 @@ def test_tombstones_modules(self): m = self.launch_bug_report_module(Tombstones) assert len(m.results) == 2 assert m.results[1]["pid"] == 3559 + + def test_dumpstate_section_is_cached_across_parallel_modules(self): + dumpstate = ( + b"header\n" + b"------ SYSTEM PROPERTIES (getprop) ------\n" + b"[persist.sys.timezone]: [Europe/Berlin]\n" + b"------------------------------------------------------------------------------\n" + b"------\n" + b"DUMP OF SERVICE package:\n" + b"package data\n" + b"------------------------------------------------------------------------------\n" + b"DUMP OF SERVICE package:\n" + b"duplicate package data\n" + b"------------------------------------------------------------------------------\n" + b"DUMP OF SERVICE adb:\n" + b"adb data\n" + b"------------------------------------------------------------------------------\n" + ) + archive = MagicMock() + archive.open.side_effect = lambda _: BytesIO(dumpstate) + shared_lock = threading.RLock() + shared_cache = {} + modules = [] + + for _ in range(4): + module = BugReportModule() + module.from_zip(archive, ["dumpState_test.log"]) + module.resource_lock = shared_lock + module.resource_cache = shared_cache + modules.append(module) + + with ThreadPoolExecutor(max_workers=4) as executor: + sections = list( + executor.map( + lambda module: module._get_dumpsys_section( + "DUMP OF SERVICE package:" + ), + modules, + ) + ) + + assert sections == ["package data"] * 4 + assert all(section is sections[0] for section in sections) + assert modules[0]._get_dumpsys_section_bytes(b"DUMP OF SERVICE adb:") == ( + b"adb data" + ) + assert modules[0]._get_system_properties_text() == ( + "[persist.sys.timezone]: [Europe/Berlin]" + ) + archive.open.assert_called_once_with("dumpState_test.log") diff --git a/tests/common/test_command.py b/tests/common/test_command.py index 865f9f16a..61c497972 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 @@ def test_store_alerts_handles_bytes(self, tmp_path): 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 @@ def test_modules_run_in_stable_topological_order(self): 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 @@ def test_custom_modules_are_filtered_before_ordering(self): ] 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 @@ def test_selected_unsupported_custom_module_does_not_run(self): 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,226 @@ def test_custom_module_dependencies_use_topological_order(self): 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() + fast_block = ( + console_output.index("--- FastLog ---"), + console_output.index("fast-a"), + console_output.index("fast-b"), + ) + slow_block = ( + console_output.index("--- SlowLog ---"), + console_output.index("slow-a"), + console_output.index("slow-b"), + ) + assert fast_block == tuple(sorted(fast_block)) + assert slow_block == tuple(sorted(slow_block)) + assert fast_block[-1] < slow_block[0] or slow_block[-1] < fast_block[0] + 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 eb10b1315..fa799d5c2 100644 --- a/tests/test_check_ios_backup.py +++ b/tests/test_check_ios_backup.py @@ -33,3 +33,11 @@ def test_check_rejects_non_backup_folder(self, tmp_path, caplog): 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 a00faea44..fd3bfc8ae 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 _record(command): 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"]}, }