Merge branch '3390-payload-plugin-interface' into develop

Issue #3390
PR #3527

monkey/infection_monkey/monkey.py

@@ -85,6 +85,7 @@
 )
 from infection_monkey.plugin.exploiter_plugin_factory import ExploiterPluginFactory
 from infection_monkey.plugin.multiprocessing_plugin_wrapper import MultiprocessingPluginWrapper
+from infection_monkey.plugin.payload_plugin_factory import PayloadPluginFactory
 from infection_monkey.propagation_credentials_repository import PropagationCredentialsRepository
 from infection_monkey.puppet import (
     PluginCompatibilityVerifier,
@@ -429,6 +430,9 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet:
                 otp_provider,
                 create_plugin,
             ),
+            AgentPluginType.PAYLOAD: PayloadPluginFactory(
+                self._agent_id, self._agent_event_publisher, create_plugin
+            ),
         }
         plugin_registry = PluginRegistry(
             operating_system,

monkey/infection_monkey/plugin/payload_plugin_factory.py

@@ -0,0 +1,27 @@
+from typing import Callable
+
+from serpentarium import SingleUsePlugin
+
+from common.event_queue import IAgentEventPublisher
+from common.types import AgentID
+
+from .i_plugin_factory import IPluginFactory
+
+
+class PayloadPluginFactory(IPluginFactory):
+    def __init__(
+        self,
+        agent_id: AgentID,
+        agent_event_publisher: IAgentEventPublisher,
+        create_plugin: Callable[..., SingleUsePlugin],
+    ):
+        self._agent_id = agent_id
+        self._agent_event_publisher = agent_event_publisher
+        self._create_plugin = create_plugin
+
+    def create(self, plugin_name: str) -> SingleUsePlugin:
+        return self._create_plugin(
+            plugin_name=plugin_name,
+            agent_id=self._agent_id,
+            agent_event_publisher=self._agent_event_publisher,
+        )

monkey/infection_monkey/puppet/puppet.py

@@ -52,7 +52,7 @@ def run_credentials_collector(
         )
         if not compatible_with_local_os:
             raise IncompatibleLocalOperatingSystemError(
-                f'The credentials collector, "{name}" is not compatible with the '
+                f'The credentials collector, "{name}", is not compatible with the '
                 "local operating system"
             )
 
@@ -104,7 +104,7 @@ def exploit_host(
         )
         if not compatible_with_local_os:
             raise IncompatibleLocalOperatingSystemError(
-                f'The exploiter, "{name}" is not compatible with the local operating system'
+                f'The exploiter, "{name}", is not compatible with the local operating system'
             )
 
         compatible_with_target_os = (
@@ -139,6 +139,16 @@ def exploit_host(
         return exploiter_result_data
 
     def run_payload(self, name: str, options: Dict, interrupt: Event):
+        compatible_with_local_os = (
+            self._plugin_compatibility_verifier.verify_local_operating_system_compatibility(
+                AgentPluginType.PAYLOAD, name
+            )
+        )
+        if not compatible_with_local_os:
+            raise IncompatibleLocalOperatingSystemError(
+                f'The payload, "{name}", is not compatible with the local operating system'
+            )
+
         payload = self._plugin_registry.get_plugin(AgentPluginType.PAYLOAD, name)
         payload.run(options, interrupt)
 

monkey/tests/unit_tests/infection_monkey/puppet/test_plugin_registry.py

@@ -128,6 +128,7 @@ def mock_plugin_factories() -> Dict[AgentPluginType, IPluginFactory]:
     return {
         AgentPluginType.EXPLOITER: MagicMock(spec=IPluginFactory),
         AgentPluginType.CREDENTIALS_COLLECTOR: MagicMock(spec=IPluginFactory),
+        AgentPluginType.PAYLOAD: MagicMock(spec=IPluginFactory),
     }
 
 
@@ -193,7 +194,7 @@ def test_get_plugin__loads_supported_plugin_types(
 
 @pytest.mark.parametrize(
     "plugin_type",
-    [AgentPluginType.FINGERPRINTER, AgentPluginType.PAYLOAD],
+    [AgentPluginType.FINGERPRINTER],
 )
 def test_get_plugin__raises_error_for_unsupported_plugin_types(
     plugin_registry: PluginRegistry, plugin_type

monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py

@@ -114,10 +114,27 @@ def test_puppet_run_multiple_payloads(puppet: Puppet):
     payload_3.run.assert_called_once()
 
 
+def test_run_payload__incompatible_local_os(
+    mock_plugin_compatibility_verifier: PluginCompatibilityVerifier, puppet: Puppet
+):
+    mock_plugin_compatibility_verifier.verify_local_operating_system_compatibility = MagicMock(  # type: ignore [assignment]  # noqa: E501
+        return_value=False
+    )
+
+    with pytest.raises(IncompatibleLocalOperatingSystemError):
+        puppet.run_payload("test", {}, threading.Event())
+
+
 def test_fingerprint_exception_handling(puppet: Puppet, mock_plugin_registry: PluginRegistry):
     mock_plugin_registry.get_plugin = MagicMock(side_effect=Exception)  # type: ignore [assignment]
     assert (
-        puppet.fingerprint("", "", PingScanData(response_received=False, os="windows"), {}, {})
+        puppet.fingerprint(
+            "",
+            "",
+            PingScanData(response_received=False, os="windows"),
+            {},  # type: ignore [arg-type]
+            {},  # type: ignore [arg-type]
+        )
         == EMPTY_FINGERPRINT
     )
 

vulture_allowlist.py

@@ -18,6 +18,7 @@
 from infection_monkey.exploit.zerologon import NetrServerPasswordSet, NetrServerPasswordSetResponse
 from infection_monkey.exploit.zerologon_utils.remote_shell import RemoteShell
 from infection_monkey.network.firewall import FirewallApp, WinAdvFirewall, WinFirewall
+from infection_monkey.plugin.payload_plugin_factory import PayloadPluginFactory
 from infection_monkey.utils import commands
 from infection_monkey.utils.decorators import request_cache
 from monkey_island.cc.deployment import Deployment
@@ -159,3 +160,6 @@
 commands.build_download_command_windows_powershell_webrequest
 
 request_cache
+
+# Remove after #3391
+PayloadPluginFactory