Merge branch '3167-mock-credential-collector' into develop

Issue #3167
PR #3247

Author: mssalvatore

envs/monkey_zoo/blackbox/test_configurations/credentials_reuse_ssh_key.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import Dict, Mapping
 
-from common.agent_configuration import AgentConfiguration, PluginConfiguration
+from common.agent_configuration import AgentConfiguration
 from common.credentials import Credentials, Password, Username
 
 from .noop import noop_test_configuration
@@ -34,10 +34,7 @@ def _add_subnets(agent_configuration: AgentConfiguration) -> AgentConfiguration:
 
 
 def _add_credentials_collectors(agent_configuration: AgentConfiguration) -> AgentConfiguration:
-    credentials_collectors = [
-        PluginConfiguration(name="SSHCollector", options={}),
-    ]
-
+    credentials_collectors: Dict[str, Mapping] = {"SSHCollector": {}}
     return add_credentials_collectors(
         agent_configuration, credentials_collectors=credentials_collectors
     )

envs/monkey_zoo/blackbox/test_configurations/depth_1_a.py

@@ -64,8 +64,9 @@ def _add_subnets(agent_configuration: AgentConfiguration) -> AgentConfiguration:
 
 
 def _add_credentials_collectors(agent_configuration: AgentConfiguration) -> AgentConfiguration:
+    credentials_collectors: Dict[str, Mapping] = {"MimikatzCollector": {}}
     return add_credentials_collectors(
-        agent_configuration, [PluginConfiguration(name="MimikatzCollector", options={})]
+        agent_configuration, credentials_collectors=credentials_collectors
     )
 
 

envs/monkey_zoo/blackbox/test_configurations/noop.py

@@ -39,7 +39,7 @@
 
 _agent_configuration = AgentConfiguration(
     keep_tunnel_open_time=0,
-    credentials_collectors=[],
+    credentials_collectors={},
     payloads={},
     propagation=_propagation_configuration,
 )

envs/monkey_zoo/blackbox/test_configurations/utils.py

@@ -45,10 +45,10 @@ def add_subnets(
 
 
 def add_credentials_collectors(
-    agent_configuration: AgentConfiguration, credentials_collectors: Sequence[PluginConfiguration]
+    agent_configuration: AgentConfiguration, credentials_collectors: Dict[str, Mapping]
 ) -> AgentConfiguration:
     agent_configuration_copy = agent_configuration.copy(deep=True)
-    agent_configuration_copy.credentials_collectors = tuple(credentials_collectors)
+    agent_configuration_copy.credentials_collectors = credentials_collectors
 
     return agent_configuration_copy
 

envs/monkey_zoo/blackbox/test_configurations/wmi_mimikatz.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import Dict, Mapping
 
-from common.agent_configuration import AgentConfiguration, PluginConfiguration
+from common.agent_configuration import AgentConfiguration
 from common.credentials import Credentials, Password, Username
 
 from .noop import noop_test_configuration
@@ -33,8 +33,9 @@ def _add_subnets(agent_configuration: AgentConfiguration) -> AgentConfiguration:
 
 
 def _add_credentials_collectors(agent_configuration: AgentConfiguration) -> AgentConfiguration:
+    credentials_collectors: Dict[str, Mapping] = {"MimikatzCollector": {}}
     return add_credentials_collectors(
-        agent_configuration, [PluginConfiguration(name="MimikatzCollector", options={})]
+        agent_configuration, credentials_collectors=credentials_collectors
     )
 
 
@@ -48,7 +49,6 @@ def _add_tcp_ports(agent_configuration: AgentConfiguration) -> AgentConfiguratio
 test_agent_configuration = _add_subnets(test_agent_configuration)
 test_agent_configuration = _add_credentials_collectors(test_agent_configuration)
 test_agent_configuration = _add_tcp_ports(test_agent_configuration)
-test_agent_configuration = _add_credentials_collectors(test_agent_configuration)
 
 CREDENTIALS = (
     Credentials(identity=Username(username="Administrator"), secret=None),

monkey/common/agent_configuration/agent_configuration.py

@@ -1,10 +1,10 @@
-from typing import Dict, Tuple
+from typing import Dict
 
 from pydantic import Field, confloat
 
 from common.base_models import MutableInfectionMonkeyBaseModel
 
-from .agent_sub_configurations import PluginConfiguration, PropagationConfiguration
+from .agent_sub_configurations import PropagationConfiguration
 
 
 class AgentConfiguration(MutableInfectionMonkeyBaseModel):
@@ -15,9 +15,11 @@ class AgentConfiguration(MutableInfectionMonkeyBaseModel):
         "seconds)",
         default=30,
     )
-    credentials_collectors: Tuple[PluginConfiguration, ...] = Field(
-        title="Credentials collectors",
-        description="Configure options for the attack’s credentials collection stage",
+    credentials_collectors: Dict[str, Dict] = Field(
+        title="Enabled credentials collectors",
+        description="Click on a credentials collector to get more information"
+        " about it. \n \u26A0 Note that using unsafe options may"
+        " result in unexpected behavior on the machine.",
     )
     payloads: Dict[str, Dict] = Field(
         title="Payloads", description="Configure payloads that Agents will execute"

monkey/common/agent_configuration/default_agent_configuration.py

@@ -1,4 +1,5 @@
 from copy import deepcopy
+from typing import Dict
 
 from . import AgentConfiguration
 from .agent_sub_configurations import (
@@ -12,11 +13,7 @@
     TCPScanConfiguration,
 )
 
-CREDENTIALS_COLLECTORS = ("MimikatzCollector", "SSHCollector")
-
-CREDENTIALS_COLLECTOR_CONFIGURATION = tuple(
-    PluginConfiguration(name=collector, options={}) for collector in CREDENTIALS_COLLECTORS
-)
+CREDENTIALS_COLLECTORS: Dict[str, Dict] = {"MimikatzCollector": {}, "SSHCollector": {}}
 
 RANSOMWARE_OPTIONS = {
     "encryption": {
@@ -93,10 +90,10 @@
 
 DEFAULT_AGENT_CONFIGURATION = AgentConfiguration(
     keep_tunnel_open_time=30,
-    credentials_collectors=CREDENTIALS_COLLECTOR_CONFIGURATION,
+    credentials_collectors=CREDENTIALS_COLLECTORS,
     payloads=PAYLOAD_CONFIGURATION,
     propagation=PROPAGATION_CONFIGURATION,
 )
 
 DEFAULT_RANSOMWARE_AGENT_CONFIGURATION = deepcopy(DEFAULT_AGENT_CONFIGURATION)
-DEFAULT_RANSOMWARE_AGENT_CONFIGURATION.credentials_collectors = tuple()
+DEFAULT_RANSOMWARE_AGENT_CONFIGURATION.credentials_collectors = {}

monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py

@@ -31,7 +31,7 @@ def __init__(self, agent_event_queue: IAgentEventQueue, agent_id: AgentID):
         self._agent_event_queue = agent_event_queue
         self._agent_id = agent_id
 
-    def collect_credentials(self, options=None) -> Sequence[Credentials]:
+    def run(self, options=None, interrupt=None) -> Sequence[Credentials]:
         logger.info("Attempting to collect windows credentials with pypykatz.")
         windows_credentials = pypykatz_handler.get_windows_creds()
 

monkey/infection_monkey/credential_collectors/ssh_collector/ssh_credential_collector.py

@@ -19,7 +19,7 @@ def __init__(self, agent_event_queue: IAgentEventQueue, agent_id: AgentID):
         self._agent_event_queue = agent_event_queue
         self._agent_id = agent_id
 
-    def collect_credentials(self, _options=None) -> Sequence[Credentials]:
+    def run(self, options=None, interrupt=None) -> Sequence[Credentials]:
         logger.info("Started scanning for SSH credentials")
         ssh_info = ssh_handler.get_ssh_info(self._agent_event_queue, self._agent_id)
         logger.info("Finished scanning for SSH credentials")

monkey/infection_monkey/i_puppet/i_credentials_collector.py

@@ -2,9 +2,10 @@
 from typing import Mapping, Optional, Sequence
 
 from common.credentials import Credentials
+from common.types import Event
 
 
 class ICredentialsCollector(ABC):
     @abstractmethod
-    def collect_credentials(self, options: Optional[Mapping]) -> Sequence[Credentials]:
+    def run(self, options: Optional[Mapping], interrupt: Event) -> Sequence[Credentials]:
         pass

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -27,32 +27,33 @@ def load_plugin(self, plugin_type: AgentPluginType, plugin_name: str, plugin: ob
         """
         Loads a plugin into the puppet
 
-        :param AgentPluginType plugin_type: The type of plugin being loaded
-        :param str plugin_name: The plugin class name
-        :param object plugin: The plugin object to load
+        :param plugin_type: The type of plugin being loaded
+        :param plugin_name: The plugin class name
+        :param plugin: The plugin object to load
         """
 
     @abc.abstractmethod
-    def run_credentials_collector(self, name: str, options: Dict) -> Sequence[Credentials]:
+    def run_credentials_collector(
+        self, name: str, options: Dict, interrupt: Event
+    ) -> Sequence[Credentials]:
         """
         Runs a credentials collector
 
-        :param str name: The name of the credentials collector to run
-        :param Dict options: A dictionary containing options that modify the behavior of the
-                             Credentials collector
+        :param name: The name of the credentials collector to run
+        :param options: A dictionary containing options that modify the behavior of the
+                        Credentials collector
+        :param interrupt: An event that can be used to interrupt the credentials collector
         :return: A sequence of Credentials that have been collected from the system
-        :rtype: Sequence[Credentials]
         """
 
     @abc.abstractmethod
     def ping(self, host: str, timeout: float) -> PingScanData:
         """
         Sends a ping (ICMP packet) to a remote host
 
-        :param str host: The domain name or IP address of a host
-        :param float timeout: The maximum amount of time (in seconds) to wait for a response
+        :param host: The domain name or IP address of a host
+        :param timeout: The maximum amount of time (in seconds) to wait for a response
         :return: The data collected by attempting to ping the target host
-        :rtype: PingScanData
         """
 
     @abc.abstractmethod
@@ -84,8 +85,7 @@ def fingerprint(
         :param name: The name of the fingerprinter to run
         :param host: The domain name or IP address of a host
         :param ping_scan_data: Data retrieved from the target host via ICMP
-        :param port_scan_data: Data retrieved from the target host via a TCP
-                                                       port scan
+        :param port_scan_data: Data retrieved from the target host via a TCP port scan
         :param options: A dictionary containing options that modify the behavior of the
                         fingerprinter
         :return: Detailed information about the target host
@@ -104,29 +104,27 @@ def exploit_host(
         """
         Runs an exploiter against a remote host
 
-        :param str name: The name of the exploiter to run
-        :param TargetHost host: A TargetHost object representing the target to exploit
-        :param int current_depth: The current propagation depth
+        :param name: The name of the exploiter to run
+        :param host: A TargetHost object representing the target to exploit
+        :param current_depth: The current propagation depth
         :param servers: List of socket addresses for victim to connect back to
-        :param Dict options: A dictionary containing options that modify the behavior of the
-                             exploiter
-        :param Event interrupt: An `Event` object that signals the exploit to stop
-                                          executing and clean itself up.
+        :param options: A dictionary containing options that modify the behavior of the exploiter
+        :param interrupt: An `Event` object that signals the exploit to stop executing and clean
+                          itself up.
         :raises IncompatibleOperatingSystemError: If an exploiter is not compatible with the target
                                                   host's operating system
-        :return: True if exploitation was successful, False otherwise
-        :rtype: ExploiterResultData
+        :return: The result of the exploit attempt
         """
 
     @abc.abstractmethod
     def run_payload(self, name: str, options: Dict, interrupt: Event):
         """
         Runs a payload
 
-        :param str name: The name of the payload to run
-        :param Dict options: A dictionary containing options that modify the behavior of the payload
-        :param Event interrupt: An `Event` object that signals the payload to stop
-                                          executing and clean itself up.
+        :param name: The name of the payload to run
+        :param options: A dictionary containing options that modify the behavior of the payload
+        :param interrupt: An `Event` object that signals the payload to stop executing and clean
+                          itself up.
         """
 
     @abc.abstractmethod

monkey/infection_monkey/master/automated_master.py

@@ -2,11 +2,10 @@
 import multiprocessing
 import time
 from ipaddress import IPv4Interface
-from typing import Any, Callable, Collection, Dict, List, Optional, Sequence
+from typing import Any, Callable, Dict, List, Optional, Sequence
 
 from egg_timer import EggTimer
 
-from common.agent_configuration import PluginConfiguration
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 from infection_monkey.i_master import IMaster
 from infection_monkey.i_puppet import IPuppet
@@ -127,7 +126,7 @@ def _run_simulation(self):
             return
 
         credentials_collector_thread = create_daemon_thread(
-            target=self._run_plugins_legacy,
+            target=self._run_plugins,
             name="CredentialsCollectorThread",
             args=(
                 config.credentials_collectors,
@@ -156,11 +155,13 @@ def _run_simulation(self):
         payload_thread.start()
         payload_thread.join()
 
-    def _collect_credentials(self, collector: PluginConfiguration):
-        credentials = self._puppet.run_credentials_collector(collector.name, collector.options)
+    def _collect_credentials(self, collector_name: str, collector_options: Dict[str, Any]):
+        credentials = self._puppet.run_credentials_collector(
+            collector_name, collector_options, self._stop
+        )
 
         if not credentials:
-            logger.debug(f"No credentials were collected by {collector}")
+            logger.debug(f"No credentials were collected by {collector_name}")
 
     def _run_payload(self, name: str, options: Dict):
         self._puppet.run_payload(name, options, self._stop)
@@ -186,26 +187,5 @@ def _run_plugins(
 
         logger.info(f"Finished running {plugin_type}s")
 
-    def _run_plugins_legacy(
-        self,
-        plugins: Collection[PluginConfiguration],
-        plugin_type: str,
-        callback: Callable[[Any], None],
-    ):
-        logger.info(f"Running {plugin_type}s")
-        logger.debug(f"Found {len(plugins)} {plugin_type}(s) to run")
-
-        interrupted_message = f"Received a stop signal, skipping remaining {plugin_type}s"
-        for p in interruptible_iter(plugins, self._stop, interrupted_message):
-            try:
-                callback(p)
-            except Exception:
-                logger.exception(
-                    f"Got unhandled exception when running {plugin_type} plugin {p}. "
-                    f"Plugin was passed to {callback}"
-                )
-
-        logger.info(f"Finished running {plugin_type}s")
-
     def cleanup(self):
         pass

monkey/infection_monkey/puppet/plugin_registry.py

@@ -1,7 +1,7 @@
 import logging
 from copy import copy
 from threading import RLock
-from typing import Any, Dict
+from typing import Dict
 
 from serpentarium import SingleUsePlugin
 
@@ -42,7 +42,16 @@ def __init__(
 
         self._lock = RLock()
 
-    def get_plugin(self, plugin_type: AgentPluginType, plugin_name: str) -> Any:
+    def get_plugin(self, plugin_type: AgentPluginType, plugin_name: str) -> SingleUsePlugin:
+        """
+        Retrieve a plugin stored in the registry. If the plugin does not exist in the registry, an
+        attempt will be made to download it from the island.
+
+        :param plugin_type: The type of plugin to get.
+        :param plugin_name: The name of the plugin to get.
+        :return: A plugin of the given type and name.
+        :raises UnknownPluginError: If the plugin is not found.
+        """
         with self._lock:
             try:
                 # Note: The MultiprocessingPluginWrapper is a MultiUsePlugin. The copy() used here

monkey/infection_monkey/puppet/puppet.py

@@ -41,11 +41,13 @@ def __init__(
     def load_plugin(self, plugin_type: AgentPluginType, plugin_name: str, plugin: object) -> None:
         self._plugin_registry.load_plugin(plugin_type, plugin_name, plugin)
 
-    def run_credentials_collector(self, name: str, options: Dict) -> Sequence[Credentials]:
+    def run_credentials_collector(
+        self, name: str, options: Dict, interrupt: Event
+    ) -> Sequence[Credentials]:
         credentials_collector = self._plugin_registry.get_plugin(
             AgentPluginType.CREDENTIALS_COLLECTOR, name
         )
-        return credentials_collector.collect_credentials(options)
+        return credentials_collector.run(options=options, interrupt=interrupt)
 
     def ping(self, host: str, timeout: float = CONNECTION_TIMEOUT) -> PingScanData:
         return network_scanning.ping(host, timeout, self._agent_event_queue, self._agent_id)