mirror of
https://github.com/mvt-project/mvt.git
synced 2026-10-31 18:24:10 +01:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
06cea62a40 |
No files matched your search
@@ -7,6 +7,7 @@ import json
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
import time
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from functools import cached_property
|
from functools import cached_property
|
||||||
from heapq import heappop, heappush
|
from heapq import heappop, heappush
|
||||||
@@ -36,6 +37,17 @@ from .utils import (
|
|||||||
from .version import MVT_VERSION
|
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:
|
class Command:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -77,6 +89,9 @@ class Command:
|
|||||||
# declaring `replaces` to the module which took their place.
|
# declaring `replaces` to the module which took their place.
|
||||||
self.module_replacements: dict[type[MVTModule], type[MVTModule]] = {}
|
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.
|
# This list will contain all executed modules.
|
||||||
# We can use this to reference e.g. self.executed[0].results.
|
# We can use this to reference e.g. self.executed[0].results.
|
||||||
self.executed: list[MVTModule] = []
|
self.executed: list[MVTModule] = []
|
||||||
@@ -190,6 +205,7 @@ class Command:
|
|||||||
"date": convert_datetime_to_iso(datetime.now()),
|
"date": convert_datetime_to_iso(datetime.now()),
|
||||||
"ioc_files": [],
|
"ioc_files": [],
|
||||||
"hashes": [],
|
"hashes": [],
|
||||||
|
"module_stats": self.module_stats,
|
||||||
}
|
}
|
||||||
|
|
||||||
for coll in self.iocs.ioc_collections:
|
for coll in self.iocs.ioc_collections:
|
||||||
@@ -766,6 +782,8 @@ class Command:
|
|||||||
except NotImplementedError:
|
except NotImplementedError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
started = time.perf_counter()
|
||||||
|
peak_before = _peak_rss_mb()
|
||||||
try:
|
try:
|
||||||
run_module(m)
|
run_module(m)
|
||||||
except EncryptedBackupError:
|
except EncryptedBackupError:
|
||||||
@@ -775,6 +793,19 @@ class Command:
|
|||||||
)
|
)
|
||||||
return
|
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)
|
self.executed.append(m)
|
||||||
executed_by_type[module] = m
|
executed_by_type[module] = m
|
||||||
self.timeline.extend(m.timeline)
|
self.timeline.extend(m.timeline)
|
||||||
|
|||||||
@@ -298,6 +298,21 @@ class TestCommand:
|
|||||||
alerts = json.loads((tmp_path / "alerts.json").read_text())
|
alerts = json.loads((tmp_path / "alerts.json").read_text())
|
||||||
assert alerts[0]["event"]["payload"] == "\\xa8\\xa9"
|
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):
|
def test_stores_collected_urls(self, tmp_path):
|
||||||
cmd = RecordingCommand(results_path=str(tmp_path))
|
cmd = RecordingCommand(results_path=str(tmp_path))
|
||||||
cmd.modules = [URLRecordingModule]
|
cmd.modules = [URLRecordingModule]
|
||||||
|
|||||||
Reference in new issue
Block a user