p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Accept list of servers in AutomatedMaster

This commit is contained in:
Ilija Lazoroski 2022-09-05 12:41:43 +02:00 committed by Mike Salvatore
parent a79b579cdc
commit 8b2107d197
2 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
monkey/infection_monkey/master/automated_master.py
View File

@ -2,7 +2,7 @@ import logging
import threading import threading
import time import time
from ipaddress import IPv4Interface from ipaddress import IPv4Interface
from typing import Any, Callable, Collection, List, Optional from typing import Any, Callable, Collection, List, Optional, Sequence
from common.agent_configuration import CustomPBAConfiguration, PluginConfiguration from common.agent_configuration import CustomPBAConfiguration, PluginConfiguration
from common.utils import Timer from common.utils import Timer
@ -35,6 +35,7 @@ class AutomatedMaster(IMaster):
def __init__( def __init__(
self, self,
current_depth: Optional[int], current_depth: Optional[int],
servers: Sequence[str],
puppet: IPuppet, puppet: IPuppet,
telemetry_messenger: ITelemetryMessenger, telemetry_messenger: ITelemetryMessenger,
victim_host_factory: VictimHostFactory, victim_host_factory: VictimHostFactory,
@ -43,6 +44,7 @@ class AutomatedMaster(IMaster):
credentials_store: IPropagationCredentialsRepository, credentials_store: IPropagationCredentialsRepository,
): ):
self._current_depth = current_depth self._current_depth = current_depth
self._servers = servers
self._puppet = puppet self._puppet = puppet
self._telemetry_messenger = telemetry_messenger self._telemetry_messenger = telemetry_messenger
self._control_channel = control_channel self._control_channel = control_channel
@ -175,7 +177,7 @@ class AutomatedMaster(IMaster):
logger.info(f"Current depth is {current_depth}") logger.info(f"Current depth is {current_depth}")
if maximum_depth_reached(config.propagation.maximum_depth, current_depth): if maximum_depth_reached(config.propagation.maximum_depth, current_depth):
self._propagator.propagate(config.propagation, current_depth, self._stop) self._propagator.propagate(config.propagation, current_depth, self._servers, self._stop)
else: else:
logger.info("Skipping propagation: maximum depth reached") logger.info("Skipping propagation: maximum depth reached")

6
monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py
View File

@ -14,7 +14,7 @@ INTERVAL = 0.001
def test_terminate_without_start(): def test_terminate_without_start():
m = AutomatedMaster(None, None, None, None, MagicMock(), [], MagicMock()) m = AutomatedMaster(None, [], None, None, None, MagicMock(), [], MagicMock())
# Test that call to terminate does not raise exception # Test that call to terminate does not raise exception
m.terminate() m.terminate()
@ -34,7 +34,7 @@ def test_stop_if_cant_get_config_from_island(monkeypatch):
monkeypatch.setattr( monkeypatch.setattr(
"infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL
) )
m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) m = AutomatedMaster(None, [], None, None, None, cc, [], MagicMock())
m.start() m.start()
assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT
@ -73,7 +73,7 @@ def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, sleep_and_return_
"infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL "infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL
) )
m = AutomatedMaster(None, None, None, None, cc, [], MagicMock()) m = AutomatedMaster(None, [], None, None, None, cc, [], MagicMock())
m.start() m.start()
assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT