CalvinBackup/mvt
1
0
Fork 0
mirror of https://github.com/mvt-project/mvt.git synced 2026-07-29 15:48:48 +02:00
Code Issues Packages Projects Releases Wiki Activity

Run extraction modules in parallel

Browse Source
This commit is contained in:
Janik Besendorf
2026-07-15 09:40:18 +02:00
parent afcfda4720
commit 9d71661f55
21 changed files with 643 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+5
View File
@@ -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
docs/development.md
+33 -1
View File
@@ -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
src/mvt/android/cli.py
+21
View File
@@ -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:
src/mvt/android/cmd_check_androidqf.py
+5
View File
@@ -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()
src/mvt/android/cmd_check_backup.py
+2
View File
@@ -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"
src/mvt/android/cmd_check_bugreport.py
+2
View File
@@ -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"
src/mvt/android/cmd_check_intrusion_logs.py
+2
View File
@@ -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"
src/mvt/android/modules/androidqf/aqf_packages.py
+3 -6
View File
@@ -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:
src/mvt/android/modules/androidqf/base.py
+8 -7
View File
@@ -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
src/mvt/android/modules/backup/base.py
+17 -16
View File
@@ -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
src/mvt/android/modules/bugreport/base.py
+17 -13
View File
@@ -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:
src/mvt/common/command.py
+199 -25
View File
@@ -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
src/mvt/common/help.py
+1
View File
@@ -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"
src/mvt/common/log.py
+41 -1
View File
@@ -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)
src/mvt/common/module.py
+16 -8
View File
@@ -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."""
src/mvt/ios/cli.py
+11
View File
@@ -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:
src/mvt/ios/cmd_check_backup.py
+2
View File
@@ -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"
src/mvt/ios/cmd_check_fs.py
+2
View File
@@ -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"
tests/common/test_command.py
+236 -5
View File
@@ -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
tests/test_check_ios_backup.py
+8
View File
@@ -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
tests/test_custom_modules.py
+12 -8
View File
@@ -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"]},
}