mirror of
https://github.com/mvt-project/mvt.git
synced 2026-10-22 15:18:26 +02:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
06cea62a40 |
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in new issue
Block a user