Agent: Change the Agent to use the modified exploiter config

monkey/infection_monkey/master/exploiter.py

@@ -2,13 +2,10 @@
 import queue
 import threading
 from queue import Queue
-from typing import Callable, Dict, Sequence
+from typing import Callable, Dict, Mapping, Sequence
 
 from common import OperatingSystem
-from common.agent_configuration.agent_sub_configurations import (
-    ExploitationConfiguration,
-    PluginConfiguration,
-)
+from common.agent_configuration.agent_sub_configurations import ExploitationConfiguration
 from common.types import Event
 from infection_monkey.custom_types import PropagationCredentials
 from infection_monkey.i_puppet import ExploiterResultData, IPuppet
@@ -22,7 +19,6 @@
 ExploiterName = str
 Callback = Callable[[ExploiterName, TargetHost, ExploiterResultData], None]
 
-
 SUPPORTED_OS = {
     "HadoopExploiter": [OperatingSystem.LINUX, OperatingSystem.WINDOWS],
     "Log4ShellExploiter": [OperatingSystem.LINUX, OperatingSystem.WINDOWS],
@@ -56,14 +52,14 @@ def exploit_hosts(
         scan_completed: threading.Event,
         stop: Event,
     ):
-        exploiters_to_run = self._process_exploiter_config(exploitation_config)
+        exploiter_configs = self._process_exploiter_config(exploitation_config)
         logger.debug(
             "Agent is configured to run the following exploiters in order: "
-            f"{', '.join([e.name for e in exploiters_to_run])}"
+            f"{', '.join([e for e in exploiter_configs])}"
         )
 
         exploit_args = (
-            exploiters_to_run,
+            exploiter_configs,
             hosts_to_exploit,
             current_depth,
             servers,
@@ -81,19 +77,19 @@ def exploit_hosts(
     @staticmethod
     def _process_exploiter_config(
         exploitation_config: ExploitationConfiguration,
-    ) -> Sequence[PluginConfiguration]:
-        extended_exploiters = []
-        for exploiter in exploitation_config.exploiters:
+    ) -> Dict[ExploiterName, Mapping]:
+        extended_configs: Dict[str, Mapping] = {}
+        for exploiter, exploiter_options in exploitation_config.exploiters.items():
             # This order allows exploiter-specific options to
             # override general options for all exploiters.
-            options = {**exploitation_config.options.__dict__, **exploiter.options}
-            extended_exploiters.append(PluginConfiguration(name=exploiter.name, options=options))
+            options = {**exploitation_config.options.__dict__, **exploiter_options}
+            extended_configs[exploiter] = options
 
-        return extended_exploiters
+        return extended_configs
 
     def _exploit_hosts_on_queue(
         self,
-        exploiters_to_run: Sequence[PluginConfiguration],
+        exploiter_configs: Dict[ExploiterName, Mapping],
         hosts_to_exploit: Queue,
         current_depth: int,
         servers: Sequence[str],
@@ -107,7 +103,7 @@ def _exploit_hosts_on_queue(
             try:
                 target_host = hosts_to_exploit.get(timeout=QUEUE_TIMEOUT)
                 self._run_all_exploiters(
-                    exploiters_to_run, target_host, current_depth, servers, results_callback, stop
+                    exploiter_configs, target_host, current_depth, servers, results_callback, stop
                 )
             except queue.Empty:
                 if _all_hosts_have_been_processed(scan_completed, hosts_to_exploit):
@@ -121,16 +117,15 @@ def _exploit_hosts_on_queue(
 
     def _run_all_exploiters(
         self,
-        exploiters_to_run: Sequence[PluginConfiguration],
+        exploiter_configs: Dict[ExploiterName, Mapping],
         target_host: TargetHost,
         current_depth: int,
         servers: Sequence[str],
         results_callback: Callback,
         stop: Event,
     ):
 
-        for exploiter in interruptible_iter(exploiters_to_run, stop):
-            exploiter_name = exploiter.name
+        for exploiter_name, exploiter_config in interruptible_iter(exploiter_configs, stop).items():
             victim_os = target_host.operating_system
 
             # This is a hack for now
@@ -144,7 +139,12 @@ def _run_all_exploiters(
                     continue
 
             exploiter_results = self._run_exploiter(
-                exploiter_name, exploiter.options, target_host, current_depth, servers, stop
+                exploiter_name,
+                exploiter_config,
+                target_host,
+                current_depth,
+                servers,
+                stop,
             )
             results_callback(exploiter_name, target_host, exploiter_results)
 
@@ -154,7 +154,7 @@ def _run_all_exploiters(
     def _run_exploiter(
         self,
         exploiter_name: str,
-        options: Dict,
+        options: Mapping,
         target_host: TargetHost,
         current_depth: int,
         servers: Sequence[str],

monkey/tests/unit_tests/infection_monkey/master/test_exploiter.py

@@ -9,11 +9,7 @@
 from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet
 
 from common import OperatingSystem
-from common.agent_configuration import (
-    AgentConfiguration,
-    ExploitationConfiguration,
-    PluginConfiguration,
-)
+from common.agent_configuration import AgentConfiguration, ExploitationConfiguration
 from infection_monkey.master import Exploiter
 from infection_monkey.model import TargetHost
 
@@ -42,12 +38,12 @@ def callback() -> Callable:
 
 @pytest.fixture
 def exploiter_config(default_agent_configuration: AgentConfiguration) -> ExploitationConfiguration:
-    exploiters = [
-        PluginConfiguration(name="MSSQLExploiter", options={"timeout": 10}),
-        PluginConfiguration(name="ZerologonExploiter", options={}),
-        PluginConfiguration(name="SSHExploiter", options={}),
-        PluginConfiguration(name="WmiExploiter", options={"timeout": 10}),
-    ]
+    exploiters = {
+        "MSSQLExploiter": {"timeout": 10},
+        "ZerologonExploiter": {},
+        "SSHExploiter": {},
+        "WmiExploiter": {"timeout": 10},
+    }
     return ExploitationConfiguration(
         options=default_agent_configuration.propagation.exploitation.options,
         exploiters=exploiters,
@@ -107,7 +103,7 @@ def inner(puppet, num_workers, hosts=hosts_to_exploit):
     return inner
 
 
-def test_exploiter(callback, hosts, hosts_to_exploit, run_exploiters):
+def test_exploiter(callback, hosts, run_exploiters):
     run_exploiters(MockPuppet(), 2)
 
     assert callback.call_count == 8