Merge branch '3167-run-credentials-collectors-as-plugins' into 3167-ui-credentials-collectors

Issue #3167
PR #3260

Author: mssalvatore

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/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,13 +155,13 @@ def _run_simulation(self):
         payload_thread.start()
         payload_thread.join()
 
-    def _collect_credentials(self, collector: PluginConfiguration):
+    def _collect_credentials(self, collector_name: str, collector_options: Dict[str, Any]):
         credentials = self._puppet.run_credentials_collector(
-            collector.name, collector.options, self._stop
+            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)
@@ -188,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/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py

@@ -28,7 +28,7 @@ def patch_pypykatz(win_creds: Sequence[WindowsCredentials], monkeypatch):
 
 def collect_credentials() -> Sequence[Credentials]:
     mock_event_queue = MagicMock(spec=IAgentEventQueue)
-    return MimikatzCredentialCollector(mock_event_queue, AGENT_ID).collect_credentials()
+    return MimikatzCredentialCollector(mock_event_queue, AGENT_ID).run()
 
 
 @pytest.mark.parametrize(
@@ -133,7 +133,7 @@ def test_mimikatz_credentials_stolen_event_published(monkeypatch):
     patch_pypykatz([], monkeypatch)
 
     mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue, AGENT_ID)
-    mimikatz_credential_collector.collect_credentials()
+    mimikatz_credential_collector.run()
 
     mock_event_queue.publish.assert_called_once()
 
@@ -147,7 +147,7 @@ def test_mimikatz_credentials_stolen_event_tags(monkeypatch):
     patch_pypykatz([], monkeypatch)
 
     mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue, AGENT_ID)
-    mimikatz_credential_collector.collect_credentials()
+    mimikatz_credential_collector.run()
 
     mock_event_queue_call_args = mock_event_queue.publish.call_args[0][0]
 
@@ -164,7 +164,7 @@ def test_mimikatz_credentials_stolen_event_stolen_credentials(monkeypatch):
     patch_pypykatz(win_creds, monkeypatch)
 
     mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue, AGENT_ID)
-    collected_credentials = mimikatz_credential_collector.collect_credentials()
+    collected_credentials = mimikatz_credential_collector.run()
 
     mock_event_queue_call_args = mock_event_queue.publish.call_args[0][0]
 

monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py

@@ -22,9 +22,7 @@ def patch_ssh_handler(ssh_creds, monkeypatch):
 )
 def test_ssh_credentials_empty_results(monkeypatch, ssh_creds):
     patch_ssh_handler(ssh_creds, monkeypatch)
-    collected = SSHCredentialCollector(
-        MagicMock(spec=IAgentEventQueue), AGENT_ID
-    ).collect_credentials()
+    collected = SSHCredentialCollector(MagicMock(spec=IAgentEventQueue), AGENT_ID).run()
     assert not collected
 
 
@@ -69,7 +67,5 @@ def test_ssh_info_result_parsing(monkeypatch):
         Credentials(identity=username3, secret=None),
         Credentials(identity=None, secret=ssh_keypair3),
     ]
-    collected = SSHCredentialCollector(
-        MagicMock(spec=IAgentEventQueue), AGENT_ID
-    ).collect_credentials()
+    collected = SSHCredentialCollector(MagicMock(spec=IAgentEventQueue), AGENT_ID).run()
     assert expected == collected