diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py index b0e974aaa1a..ef8e2de5fc3 100644 --- a/monkey/infection_monkey/i_puppet/__init__.py +++ b/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 diff --git a/monkey/infection_monkey/i_puppet/payload_result.py b/monkey/infection_monkey/i_puppet/payload_result.py new file mode 100644 index 00000000000..d144de0269f --- /dev/null +++ b/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 = "" diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 0141193da61..7542403d37b 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/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 diff --git a/monkey/tests/data_for_tests/payload_plugin/MockPayload-payload.tar b/monkey/tests/data_for_tests/payload_plugin/MockPayload-payload.tar new file mode 100644 index 00000000000..ff6953e2b53 Binary files /dev/null and b/monkey/tests/data_for_tests/payload_plugin/MockPayload-payload.tar differ diff --git a/monkey/tests/data_for_tests/payload_plugin/build.sh b/monkey/tests/data_for_tests/payload_plugin/build.sh new file mode 100644 index 00000000000..52ffc9b4b81 --- /dev/null +++ b/monkey/tests/data_for_tests/payload_plugin/build.sh @@ -0,0 +1,48 @@ +#!/bin/sh + +# Build the plugin package +# Usage: ./build.sh + +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" diff --git a/monkey/tests/data_for_tests/payload_plugin/config-schema.json b/monkey/tests/data_for_tests/payload_plugin/config-schema.json new file mode 100644 index 00000000000..c56b3762d7c --- /dev/null +++ b/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 + } + } +} diff --git a/monkey/tests/data_for_tests/payload_plugin/manifest.yaml b/monkey/tests/data_for_tests/payload_plugin/manifest.yaml new file mode 100644 index 00000000000..234f76a81bf --- /dev/null +++ b/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 diff --git a/monkey/tests/data_for_tests/payload_plugin/src/plugin.py b/monkey/tests/data_for_tests/payload_plugin/src/plugin.py new file mode 100644 index 00000000000..90fe4056fe8 --- /dev/null +++ b/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 diff --git a/monkey/tests/data_for_tests/payload_plugin/src/vendor/mock_dependency.py b/monkey/tests/data_for_tests/payload_plugin/src/vendor/mock_dependency.py new file mode 100644 index 00000000000..5becc17c04a --- /dev/null +++ b/monkey/tests/data_for_tests/payload_plugin/src/vendor/mock_dependency.py @@ -0,0 +1 @@ +__version__ = "1.0.0" diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 1277b3ede7e..72b3b8b37dd 100644 --- a/vulture_allowlist.py +++ b/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