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

Issue #3390
PR #3531

monkey/infection_monkey/i_puppet/__init__.py

@@ -10,4 +10,5 @@
     IncompatibleTargetOperatingSystemError,
 )
 from .i_fingerprinter import IFingerprinter
+from .payload_result import PayloadResult
 from .target_host import TargetHost, TargetHostPorts, PortScanDataDict

monkey/infection_monkey/i_puppet/payload_result.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class PayloadResult:
+    run_success: bool = False
+    error_message: str = ""

monkey/infection_monkey/puppet/puppet.py

@@ -150,7 +150,7 @@ def run_payload(self, name: str, options: Dict, interrupt: Event):
             )
 
         payload = self._plugin_registry.get_plugin(AgentPluginType.PAYLOAD, name)
-        payload.run(options, interrupt)
+        payload.run(options=options, interrupt=interrupt)
 
     def cleanup(self) -> None:
         pass

monkey/tests/data_for_tests/payload_plugin/build.sh

@@ -0,0 +1,48 @@
+#!/bin/sh
+
+# Build the plugin package
+# Usage: ./build.sh <version>
+
+DEFAULT_DEPENDENCY_VERSION=1.0.0
+MANIFEST_FILENAME=manifest.yaml
+SCHEMA_FILENAME=config-schema.json
+DEPENDENCY_FILE="src/vendor/mock_dependency.py"
+ROOT="$( cd "$( dirname "$0" )" && pwd )"
+
+get_value_from_key() {
+    _file="$1"
+    _key="$2"
+    _value=$(grep -Po "(?<=^${_key}:).*" "$_file")
+    if [ -z "$_value" ]; then
+        echo "Error: Plugin '$_key' not found."
+        exit 1
+    else
+        echo "$_value"
+    fi
+}
+
+lower() {
+    echo "$1" | tr "[:upper:]" "[:lower:]"
+}
+
+# Generate the dependency
+version=$DEFAULT_DEPENDENCY_VERSION
+if [ "$1" ]; then
+    version=$1
+fi
+echo "__version__ = \"${version}\"" > "$ROOT/$DEPENDENCY_FILE"
+
+
+# Package everything up
+cd "$ROOT/src" || exit 1
+tar -czf $ROOT/source.tar.gz plugin.py vendor/
+cd "$ROOT" || exit 1
+
+
+# xargs strips leading whitespace
+name=$(get_value_from_key $MANIFEST_FILENAME name | xargs)
+type=$(lower "$(get_value_from_key $MANIFEST_FILENAME plugin_type | xargs)")
+
+plugin_filename="${name}-${type}.tar"
+tar -cf "$ROOT/$plugin_filename" $MANIFEST_FILENAME $SCHEMA_FILENAME source.tar.gz
+rm "$ROOT/source.tar.gz"

monkey/tests/data_for_tests/payload_plugin/config-schema.json

@@ -0,0 +1,19 @@
+{
+    "title": "Mock payload",
+    "description": "Configuration settings for mock payload.",
+    "type": "object",
+    "properties": {
+         "random_boolean": {
+            "title": "Random boolean",
+            "description": "A random boolean field for testing",
+            "type": "boolean",
+            "default": true
+        },
+        "sleep_duration": {
+            "title": "Sleep duration",
+            "description": "Duration in seconds for which the plugin should sleep",
+            "type": "number",
+            "default": 0
+        }
+    }
+}

monkey/tests/data_for_tests/payload_plugin/manifest.yaml

@@ -0,0 +1,13 @@
+name: MockPayload
+plugin_type: Payload
+supported_operating_systems:
+  - linux
+  - windows
+target_operating_systems:
+  - windows
+  - linux
+title: Mock Payload plugin
+description: A payload plugin for testing purposes
+version: 1.0.0
+safe: true
+link_to_documentation: https://www.akamai.com/infectionmonkey

monkey/tests/data_for_tests/payload_plugin/src/plugin.py

@@ -0,0 +1,72 @@
+import logging
+import time
+from pathlib import PurePosixPath
+from threading import Event, current_thread
+from typing import Any, Dict
+
+import mock_dependency
+
+from common.agent_events import AgentEventTag, FileEncryptionEvent
+from common.event_queue import IAgentEventPublisher
+from common.types import AgentID
+from infection_monkey.i_puppet import PayloadResult
+from infection_monkey.utils.threading import interruptible_iter
+
+logger = logging.getLogger(__name__)
+
+
+class Plugin:
+    def __init__(
+        self,
+        *,
+        plugin_name="",
+        agent_id: AgentID,
+        agent_event_publisher: IAgentEventPublisher,
+        **kwargs,
+    ):
+        self._agent_id = agent_id
+        self._agent_event_publisher = agent_event_publisher
+
+    def run(
+        self,
+        *,
+        options: Dict[str, Any],
+        interrupt: Event,
+        **kwargs,
+    ):
+        logger.info(f"Main thread name {current_thread().name}")
+        logger.info(f"Mock dependency package version: {mock_dependency.__version__}")
+
+        Plugin._log_options(options)
+        Plugin._sleep(options.get("sleep_duration", 0), interrupt)
+
+        return self._run_payload(options)
+
+    @staticmethod
+    def _log_options(options: Dict[str, Any]):
+        logger.info("Plugin options:")
+
+        random_boolean = options.get("random_boolean", None)
+        logger.info(f"Random boolean: {random_boolean}")
+
+    @staticmethod
+    def _sleep(duration: float, interrupt: Event):
+        logger.info(f"Sleeping for {duration} seconds")
+        for time_passed in interruptible_iter(range(int(duration)), interrupt):
+            logger.info(f"Passed {time_passed} seconds")
+            time.sleep(1)
+
+    def _run_payload(self, options: Dict[str, Any]) -> PayloadResult:
+        payload_result = PayloadResult(run_success=True)
+
+        self._agent_event_publisher.publish(
+            FileEncryptionEvent(
+                source=self._agent_id,
+                file_path=PurePosixPath("/home/ubuntu/encrypted.txt"),
+                success=True,
+                error_message="error",
+                tags=frozenset({AgentEventTag("payload-tag")}),
+            )
+        )
+
+        return payload_result

monkey/tests/data_for_tests/payload_plugin/src/vendor/mock_dependency.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

vulture_allowlist.py

@@ -17,6 +17,7 @@
 from infection_monkey.exploit.tools import secret_type_filter
 from infection_monkey.exploit.zerologon import NetrServerPasswordSet, NetrServerPasswordSetResponse
 from infection_monkey.exploit.zerologon_utils.remote_shell import RemoteShell
+from infection_monkey.i_puppet import PayloadResult
 from infection_monkey.network.firewall import FirewallApp, WinAdvFirewall, WinFirewall
 from infection_monkey.plugin.payload_plugin_factory import PayloadPluginFactory
 from infection_monkey.utils import commands
@@ -163,3 +164,5 @@
 
 # Remove after #3391
 PayloadPluginFactory
+PayloadResult
+PayloadResult.run_success