diff --git a/src/mvt/common/command.py b/src/mvt/common/command.py index 6848e8c0..cb0e0f37 100644 --- a/src/mvt/common/command.py +++ b/src/mvt/common/command.py @@ -7,6 +7,7 @@ import json import logging import os import sys +import time from datetime import datetime from functools import cached_property from heapq import heappop, heappush @@ -36,6 +37,17 @@ from .utils import ( from .version import MVT_VERSION +def _peak_rss_mb() -> float: + """The process's peak resident set size so far, in MB (0 where unknown).""" + try: + import resource + except ImportError: + return 0.0 + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # ru_maxrss is in bytes on macOS and in kilobytes elsewhere. + return peak / (1024 * 1024 if sys.platform == "darwin" else 1024) + + class Command: def __init__( self, @@ -77,6 +89,9 @@ class Command: # declaring `replaces` to the module which took their place. self.module_replacements: dict[type[MVTModule], type[MVTModule]] = {} + # Wall time and peak-RSS growth of every module run, in run order. + self.module_stats: list[dict[str, Any]] = [] + # This list will contain all executed modules. # We can use this to reference e.g. self.executed[0].results. self.executed: list[MVTModule] = [] @@ -190,6 +205,7 @@ class Command: "date": convert_datetime_to_iso(datetime.now()), "ioc_files": [], "hashes": [], + "module_stats": self.module_stats, } for coll in self.iocs.ioc_collections: @@ -766,6 +782,8 @@ class Command: except NotImplementedError: pass + started = time.perf_counter() + peak_before = _peak_rss_mb() try: run_module(m) except EncryptedBackupError: @@ -775,6 +793,19 @@ class Command: ) return + stats = { + "module": module.__name__, + "seconds": round(time.perf_counter() - started, 2), + "peak_rss_growth_mb": round(_peak_rss_mb() - peak_before), + } + self.module_stats.append(stats) + self.log.info( + "Module %s ran in %.1fs and raised the peak memory by %d MB", + stats["module"], + stats["seconds"], + stats["peak_rss_growth_mb"], + ) + self.executed.append(m) executed_by_type[module] = m self.timeline.extend(m.timeline) diff --git a/tests/common/test_command.py b/tests/common/test_command.py index 357d3b7d..2c96d4a3 100644 --- a/tests/common/test_command.py +++ b/tests/common/test_command.py @@ -298,6 +298,21 @@ class TestCommand: alerts = json.loads((tmp_path / "alerts.json").read_text()) assert alerts[0]["event"]["payload"] == "\\xa8\\xa9" + def test_records_time_and_memory_per_module(self, tmp_path): + cmd = RecordingCommand(results_path=str(tmp_path)) + cmd.modules = [FirstModule, SecondModule] + + cmd.run() + + info = json.loads((tmp_path / "info.json").read_text()) + assert [stat["module"] for stat in info["module_stats"]] == [ + "FirstModule", + "SecondModule", + ] + for stat in info["module_stats"]: + assert stat["seconds"] >= 0 + assert stat["peak_rss_growth_mb"] >= 0 + def test_stores_collected_urls(self, tmp_path): cmd = RecordingCommand(results_path=str(tmp_path)) cmd.modules = [URLRecordingModule]