mirror of
https://github.com/mvt-project/mvt.git
synced 2026-10-27 16:26:23 +01:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6a8fe14a82 | ||
|
|
529e7ff4dd | ||
|
|
3549002ec2 | ||
|
|
4530d496d3 |
No files matched your search
@@ -21,6 +21,24 @@ make check
|
|||||||
|
|
||||||
Run these tests before making new commits or opening pull requests.
|
Run these tests before making new commits or opening pull requests.
|
||||||
|
|
||||||
|
## Module dependencies
|
||||||
|
|
||||||
|
Modules can require other modules to run first by declaring their classes in
|
||||||
|
`dependencies`. The command runner uses a stable topological ordering, so the
|
||||||
|
existing module list order is preserved wherever dependency constraints allow.
|
||||||
|
|
||||||
|
```python
|
||||||
|
class DependentModule(MVTModule):
|
||||||
|
dependencies = (PrerequisiteModule,)
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
prerequisite_results = self.get_dependency_results(PrerequisiteModule)
|
||||||
|
```
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
## Profiling
|
## Profiling
|
||||||
|
|
||||||
Some MVT modules extract and process significant amounts of data during the analysis process or while checking results against known indicators. Care must be
|
Some MVT modules extract and process significant amounts of data during the analysis process or while checking results against known indicators. Care must be
|
||||||
|
|||||||
+82
-11
@@ -8,6 +8,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
from heapq import heappop, heappush
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
|
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
@@ -18,6 +19,7 @@ from .alerts import AlertLevel, AlertStore
|
|||||||
from .config import settings
|
from .config import settings
|
||||||
from .indicators import Indicators
|
from .indicators import Indicators
|
||||||
from .module import EncryptedBackupError, MVTModule, run_module, save_timeline
|
from .module import EncryptedBackupError, MVTModule, run_module, save_timeline
|
||||||
|
from .module_types import ModuleTimeline
|
||||||
from .utils import (
|
from .utils import (
|
||||||
CustomJSONEncoder,
|
CustomJSONEncoder,
|
||||||
convert_datetime_to_iso,
|
convert_datetime_to_iso,
|
||||||
@@ -44,7 +46,7 @@ class Command:
|
|||||||
disable_indicator_check: bool = False,
|
disable_indicator_check: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.name = ""
|
self.name = ""
|
||||||
self.modules: list[Any] = []
|
self.modules: list[type[MVTModule]] = []
|
||||||
|
|
||||||
self.target_path = target_path
|
self.target_path = target_path
|
||||||
self.results_path = results_path
|
self.results_path = results_path
|
||||||
@@ -63,10 +65,10 @@ class Command:
|
|||||||
|
|
||||||
# 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[Any] = []
|
self.executed: list[MVTModule] = []
|
||||||
self.hashes = hashes
|
self.hashes = hashes
|
||||||
self.hash_values: list[dict[str, Any]] = []
|
self.hash_values: list[dict[str, Any]] = []
|
||||||
self.timeline: list[dict[str, Any]] = []
|
self.timeline: ModuleTimeline = []
|
||||||
|
|
||||||
# Load IOCs
|
# Load IOCs
|
||||||
self._create_storage()
|
self._create_storage()
|
||||||
@@ -258,20 +260,84 @@ class Command:
|
|||||||
console.print("")
|
console.print("")
|
||||||
console.print(panel)
|
console.print(panel)
|
||||||
|
|
||||||
|
def _ordered_modules(self) -> Optional[list[type[MVTModule]]]:
|
||||||
|
"""Return enabled modules in stable topological order."""
|
||||||
|
module_indexes = {module: index for index, module in enumerate(self.modules)}
|
||||||
|
|
||||||
|
if self.module_name:
|
||||||
|
selected = [
|
||||||
|
module for module in self.modules if module.__name__ == self.module_name
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
selected = [module for module in self.modules if module.enabled]
|
||||||
|
|
||||||
|
required = set(selected)
|
||||||
|
pending = list(selected)
|
||||||
|
while pending:
|
||||||
|
module = pending.pop()
|
||||||
|
for dependency in module.dependencies:
|
||||||
|
if dependency not in module_indexes:
|
||||||
|
self.log.warning(
|
||||||
|
"Module %s depends on unavailable module %s. "
|
||||||
|
"No modules will be run.",
|
||||||
|
module.__name__,
|
||||||
|
dependency.__name__,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
if dependency not in required:
|
||||||
|
required.add(dependency)
|
||||||
|
pending.append(dependency)
|
||||||
|
|
||||||
|
dependents: dict[type[MVTModule], list[type[MVTModule]]] = {
|
||||||
|
module: [] for module in required
|
||||||
|
}
|
||||||
|
indegree = {module: 0 for module in required}
|
||||||
|
for module in required:
|
||||||
|
for dependency in module.dependencies:
|
||||||
|
if dependency not in required:
|
||||||
|
continue
|
||||||
|
dependents[dependency].append(module)
|
||||||
|
indegree[module] += 1
|
||||||
|
|
||||||
|
ready: list[tuple[int, type[MVTModule]]] = []
|
||||||
|
for module, count in indegree.items():
|
||||||
|
if count == 0:
|
||||||
|
heappush(ready, (module_indexes[module], module))
|
||||||
|
|
||||||
|
ordered = []
|
||||||
|
while ready:
|
||||||
|
_, module = heappop(ready)
|
||||||
|
ordered.append(module)
|
||||||
|
for dependent in dependents[module]:
|
||||||
|
indegree[dependent] -= 1
|
||||||
|
if indegree[dependent] == 0:
|
||||||
|
heappush(ready, (module_indexes[dependent], dependent))
|
||||||
|
|
||||||
|
if len(ordered) != len(required):
|
||||||
|
cyclic_modules = sorted(
|
||||||
|
(module.__name__ for module, count in indegree.items() if count > 0)
|
||||||
|
)
|
||||||
|
self.log.warning(
|
||||||
|
"Circular module dependency detected involving: %s. "
|
||||||
|
"No modules will be run.",
|
||||||
|
", ".join(cyclic_modules),
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
return ordered
|
||||||
|
|
||||||
def run(self) -> None:
|
def run(self) -> None:
|
||||||
|
ordered_modules = self._ordered_modules()
|
||||||
|
if ordered_modules is None:
|
||||||
|
return
|
||||||
|
|
||||||
try:
|
try:
|
||||||
self.init()
|
self.init()
|
||||||
except NotImplementedError:
|
except NotImplementedError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
for module in self.modules:
|
executed_by_type: dict[type[MVTModule], MVTModule] = {}
|
||||||
if self.module_name and module.__name__ != self.module_name:
|
for module in ordered_modules:
|
||||||
continue
|
|
||||||
|
|
||||||
if not module.enabled and not (
|
|
||||||
self.module_name and module.__name__ == self.module_name
|
|
||||||
):
|
|
||||||
continue
|
|
||||||
|
|
||||||
module_logger = logging.getLogger(module.__module__)
|
module_logger = logging.getLogger(module.__module__)
|
||||||
|
|
||||||
@@ -281,6 +347,10 @@ class Command:
|
|||||||
module_options=self.module_options,
|
module_options=self.module_options,
|
||||||
log=module_logger,
|
log=module_logger,
|
||||||
)
|
)
|
||||||
|
m.dependency_modules = {
|
||||||
|
dependency: executed_by_type[dependency]
|
||||||
|
for dependency in module.dependencies
|
||||||
|
}
|
||||||
|
|
||||||
if self.iocs.total_ioc_count:
|
if self.iocs.total_ioc_count:
|
||||||
m.indicators = self.iocs
|
m.indicators = self.iocs
|
||||||
@@ -304,6 +374,7 @@ class Command:
|
|||||||
return
|
return
|
||||||
|
|
||||||
self.executed.append(m)
|
self.executed.append(m)
|
||||||
|
executed_by_type[module] = m
|
||||||
self.timeline.extend(m.timeline)
|
self.timeline.extend(m.timeline)
|
||||||
self.alertstore.extend(m.alertstore.alerts)
|
self.alertstore.extend(m.alertstore.alerts)
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
from dataclasses import asdict, is_dataclass
|
from dataclasses import asdict, is_dataclass
|
||||||
from typing import Any, Dict, Optional
|
from typing import Any, Dict, Optional, Sequence
|
||||||
|
|
||||||
from .alerts import AlertStore
|
from .alerts import AlertStore
|
||||||
from .indicators import Indicators
|
from .indicators import Indicators
|
||||||
@@ -43,6 +43,7 @@ class MVTModule:
|
|||||||
|
|
||||||
enabled: bool = True
|
enabled: bool = True
|
||||||
slug: Optional[str] = None
|
slug: Optional[str] = None
|
||||||
|
dependencies: Sequence[type["MVTModule"]] = ()
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -71,6 +72,7 @@ class MVTModule:
|
|||||||
self.file_path: Optional[str] = file_path
|
self.file_path: Optional[str] = file_path
|
||||||
self.target_path: Optional[str] = target_path
|
self.target_path: Optional[str] = target_path
|
||||||
self.results_path: Optional[str] = results_path
|
self.results_path: Optional[str] = results_path
|
||||||
|
self.serial: Optional[str] = None
|
||||||
self.module_options: Dict[str, Any] = module_options if module_options else {}
|
self.module_options: Dict[str, Any] = module_options if module_options else {}
|
||||||
|
|
||||||
self.log = log
|
self.log = log
|
||||||
@@ -79,6 +81,13 @@ class MVTModule:
|
|||||||
|
|
||||||
self.results: ModuleResults = results if results is not None else []
|
self.results: ModuleResults = results if results is not None else []
|
||||||
self.timeline: ModuleTimeline = []
|
self.timeline: ModuleTimeline = []
|
||||||
|
self.dependency_modules: Dict[type["MVTModule"], "MVTModule"] = {}
|
||||||
|
|
||||||
|
def get_dependency_results(
|
||||||
|
self, module_class: type["MVTModule"]
|
||||||
|
) -> ModuleResults:
|
||||||
|
"""Return the results produced by a prerequisite module."""
|
||||||
|
return self.dependency_modules[module_class].results
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_json(cls, json_path: str, log: logging.Logger):
|
def from_json(cls, json_path: str, log: logging.Logger):
|
||||||
|
|||||||
@@ -4,11 +4,59 @@
|
|||||||
# https://license.mvt.re/1.1/
|
# https://license.mvt.re/1.1/
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
|
|
||||||
from mvt.common.command import Command
|
from mvt.common.command import Command
|
||||||
|
from mvt.common.module import MVTModule
|
||||||
|
|
||||||
|
|
||||||
|
class RecordingModule(MVTModule):
|
||||||
|
run_order: list[str] = []
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
self.run_order.append(self.__class__.__name__)
|
||||||
|
|
||||||
|
def check_indicators(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class FirstModule(RecordingModule):
|
||||||
|
def run(self):
|
||||||
|
super().run()
|
||||||
|
self.results = ["first"]
|
||||||
|
|
||||||
|
|
||||||
|
class SecondModule(RecordingModule):
|
||||||
|
dependencies = (FirstModule,)
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
super().run()
|
||||||
|
self.results = self.get_dependency_results(FirstModule) + ["second"]
|
||||||
|
|
||||||
|
|
||||||
|
class ThirdModule(RecordingModule):
|
||||||
|
dependencies = (SecondModule,)
|
||||||
|
|
||||||
|
|
||||||
|
class IndependentModule(RecordingModule):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class RecordingCommand(Command):
|
||||||
|
def init(self):
|
||||||
|
self.initialized = True
|
||||||
|
|
||||||
|
def module_init(self, module):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def finish(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
class TestCommand:
|
class TestCommand:
|
||||||
|
def setup_method(self):
|
||||||
|
RecordingModule.run_order = []
|
||||||
|
|
||||||
def test_store_alerts_handles_bytes(self, tmp_path):
|
def test_store_alerts_handles_bytes(self, tmp_path):
|
||||||
cmd = Command(results_path=str(tmp_path))
|
cmd = Command(results_path=str(tmp_path))
|
||||||
cmd.alertstore.medium(
|
cmd.alertstore.medium(
|
||||||
@@ -21,3 +69,66 @@ 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_modules_run_in_stable_topological_order(self):
|
||||||
|
cmd = RecordingCommand()
|
||||||
|
cmd.modules = [ThirdModule, IndependentModule, SecondModule, FirstModule]
|
||||||
|
|
||||||
|
cmd.run()
|
||||||
|
|
||||||
|
assert RecordingModule.run_order == [
|
||||||
|
"IndependentModule",
|
||||||
|
"FirstModule",
|
||||||
|
"SecondModule",
|
||||||
|
"ThirdModule",
|
||||||
|
]
|
||||||
|
second = next(module for module in cmd.executed if isinstance(module, SecondModule))
|
||||||
|
assert second.results == ["first", "second"]
|
||||||
|
|
||||||
|
def test_selected_module_runs_transitive_dependencies(self):
|
||||||
|
cmd = RecordingCommand(module_name="ThirdModule")
|
||||||
|
cmd.modules = [ThirdModule, SecondModule, FirstModule, IndependentModule]
|
||||||
|
|
||||||
|
cmd.run()
|
||||||
|
|
||||||
|
assert RecordingModule.run_order == [
|
||||||
|
"FirstModule",
|
||||||
|
"SecondModule",
|
||||||
|
"ThirdModule",
|
||||||
|
]
|
||||||
|
|
||||||
|
def test_circular_dependency_warns_and_stops(self, caplog):
|
||||||
|
class CircularOne(RecordingModule):
|
||||||
|
pass
|
||||||
|
|
||||||
|
class CircularTwo(RecordingModule):
|
||||||
|
dependencies = (CircularOne,)
|
||||||
|
|
||||||
|
CircularOne.dependencies = (CircularTwo,)
|
||||||
|
|
||||||
|
cmd = RecordingCommand()
|
||||||
|
cmd.modules = [CircularOne, CircularTwo]
|
||||||
|
|
||||||
|
with caplog.at_level(logging.WARNING):
|
||||||
|
cmd.run()
|
||||||
|
|
||||||
|
assert RecordingModule.run_order == []
|
||||||
|
assert not hasattr(cmd, "initialized")
|
||||||
|
assert "Circular module dependency detected" in caplog.text
|
||||||
|
|
||||||
|
def test_unavailable_dependency_warns_and_stops(self, caplog):
|
||||||
|
class UnavailableModule(RecordingModule):
|
||||||
|
pass
|
||||||
|
|
||||||
|
class DependentModule(RecordingModule):
|
||||||
|
dependencies = (UnavailableModule,)
|
||||||
|
|
||||||
|
cmd = RecordingCommand()
|
||||||
|
cmd.modules = [DependentModule]
|
||||||
|
|
||||||
|
with caplog.at_level(logging.WARNING):
|
||||||
|
cmd.run()
|
||||||
|
|
||||||
|
assert RecordingModule.run_order == []
|
||||||
|
assert not hasattr(cmd, "initialized")
|
||||||
|
assert "depends on unavailable module UnavailableModule" in caplog.text
|
||||||
Reference in new issue
Block a user