forked from p15670423/monkey
Agent: Implement _run_simulation() that calls stubbed methods
This commit is contained in:
parent
a2bba6a025
commit
9809fc2a41
|
@ -1,6 +1,7 @@
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
from infection_monkey.i_control_channel import IControlChannel
|
from infection_monkey.i_control_channel import IControlChannel
|
||||||
from infection_monkey.i_master import IMaster
|
from infection_monkey.i_master import IMaster
|
||||||
|
@ -8,7 +9,7 @@ from infection_monkey.i_puppet import IPuppet
|
||||||
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
|
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
|
||||||
|
|
||||||
CHECK_FOR_STOP_INTERVAL_SEC = 5
|
CHECK_FOR_STOP_INTERVAL_SEC = 5
|
||||||
SHUTDOWN_TIMEOUT = 2
|
SHUTDOWN_TIMEOUT = 5
|
||||||
|
|
||||||
logger = logging.getLogger()
|
logger = logging.getLogger()
|
||||||
|
|
||||||
|
@ -63,11 +64,57 @@ class AutomatedMaster(IMaster):
|
||||||
logger.warn("Forcefully killing the simulation")
|
logger.warn("Forcefully killing the simulation")
|
||||||
|
|
||||||
def _run_simulation(self):
|
def _run_simulation(self):
|
||||||
|
config = self._control_channel.get_config()
|
||||||
|
|
||||||
|
system_info_collector_thread = threading.Thread(
|
||||||
|
target=self._collect_system_info,
|
||||||
|
args=(config["system_info_collector_classes"],),
|
||||||
|
daemon=True,
|
||||||
|
)
|
||||||
|
pba_thread = threading.Thread(
|
||||||
|
target=self._run_pbas, args=(config["post_breach_actions"],), daemon=True
|
||||||
|
)
|
||||||
|
|
||||||
|
system_info_collector_thread.start()
|
||||||
|
pba_thread.start()
|
||||||
|
|
||||||
|
# Future stages of the simulation require the output of the system info collectors. Nothing
|
||||||
|
# requires the output of PBAs, so we don't need to join on that thread.
|
||||||
|
system_info_collector_thread.join()
|
||||||
|
|
||||||
|
if self._can_propagate():
|
||||||
|
propagation_thread = threading.Thread(
|
||||||
|
target=self._propagate, args=(config,), daemon=True
|
||||||
|
)
|
||||||
|
propagation_thread.start()
|
||||||
|
propagation_thread.join()
|
||||||
|
|
||||||
|
payload_thread = threading.Thread(
|
||||||
|
target=self._run_payloads, args=(config["payloads"],), daemon=True
|
||||||
|
)
|
||||||
|
payload_thread.start()
|
||||||
|
payload_thread.join()
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
time.sleep(30)
|
time.sleep(2)
|
||||||
logger.debug("Simulation thread is finished sleeping")
|
logger.debug("Simulation thread is finished sleeping")
|
||||||
if self._stop.is_set():
|
if self._stop.is_set():
|
||||||
break
|
break
|
||||||
|
|
||||||
|
def _collect_system_info(self, enabled_collectors: List[str]):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _run_pbas(self, enabled_pbas: List[str]):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _can_propagate(self):
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _propagate(self, config: Dict):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _run_payloads(self, enabled_payloads: Dict[str, Dict]):
|
||||||
|
pass
|
||||||
|
|
||||||
def cleanup(self):
|
def cleanup(self):
|
||||||
pass
|
pass
|
||||||
|
|
Loading…
Reference in New Issue