diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 0fec20edb11..b8151ec5edd 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -461,7 +461,9 @@ def _build_puppet(self, operating_system: OperatingSystem) -> IPuppet: ) puppet.load_plugin( - AgentPluginType.PAYLOAD, "ransomware", RansomwarePayload(self._agent_event_queue) + AgentPluginType.PAYLOAD, + "ransomware", + RansomwarePayload(self._agent_event_queue, self._agent_id), ) return puppet diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index 0531b75fcaa..503f5e092bd 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -6,7 +6,7 @@ from common.agent_events import FileEncryptionEvent from common.event_queue import IAgentEventQueue from common.tags import T1486_ATTACK_TECHNIQUE_TAG -from infection_monkey.utils.ids import get_agent_id +from common.types import AgentID from infection_monkey.utils.threading import interruptible_function, interruptible_iter from .consts import README_FILE_NAME, README_SRC @@ -26,6 +26,7 @@ def __init__( select_files: Callable[[Path], Iterable[Path]], leave_readme: Callable[[Path, Path], None], agent_event_queue: IAgentEventQueue, + agent_id: AgentID, ): self._config = config @@ -33,6 +34,7 @@ def __init__( self._select_files = select_files self._leave_readme = leave_readme self._agent_event_queue = agent_event_queue + self._agent_id = agent_id self._target_directory = self._config.target_directory self._readme_file_path = ( @@ -91,7 +93,7 @@ def _encrypt_files(self, files_to_encrypt: Iterable[Path], interrupt: threading. def _publish_file_encryption_event(self, filepath: Path, success: bool, error: str): file_encryption_event = FileEncryptionEvent( - source=get_agent_id(), + source=self._agent_id, file_path=filepath, success=success, error_message=error, diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py index 8f2f28f1618..ecaf94ad282 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -2,6 +2,7 @@ from pprint import pformat from common.event_queue import IAgentEventQueue +from common.types import AgentID from infection_monkey.utils.bit_manipulators import flip_bits from . import readme_dropper @@ -19,6 +20,7 @@ def build_ransomware( options: dict, agent_event_queue: IAgentEventQueue, + agent_id: AgentID, ): logger.debug(f"Ransomware configuration:\n{pformat(options)}") ransomware_options = RansomwareOptions(options) @@ -28,7 +30,7 @@ def build_ransomware( leave_readme = _build_leave_readme() return Ransomware( - ransomware_options, file_encryptor, file_selector, leave_readme, agent_event_queue + ransomware_options, file_encryptor, file_selector, leave_readme, agent_event_queue, agent_id ) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware_payload.py index 8635dedffb2..e787f59fdec 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_payload.py @@ -1,16 +1,19 @@ from typing import Dict from common.event_queue import IAgentEventQueue -from common.types import Event +from common.types import AgentID, Event from infection_monkey.payload.i_payload import IPayload from . import ransomware_builder class RansomwarePayload(IPayload): - def __init__(self, agent_event_queue: IAgentEventQueue): + def __init__(self, agent_event_queue: IAgentEventQueue, agent_id: AgentID): self._agent_event_queue = agent_event_queue + self._agent_id = agent_id def run(self, options: Dict, interrupt: Event): - ransomware = ransomware_builder.build_ransomware(options, self._agent_event_queue) + ransomware = ransomware_builder.build_ransomware( + options, self._agent_event_queue, self._agent_id + ) ransomware.run(interrupt) diff --git a/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py b/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py index 48a3eae707b..bfe0f507028 100644 --- a/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py +++ b/monkey/tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py @@ -7,6 +7,9 @@ import infection_monkey.payload.ransomware.ransomware_builder as ransomware_builder from common.agent_configuration.default_agent_configuration import RANSOMWARE_OPTIONS from common.event_queue import IAgentEventQueue +from common.types import AgentID + +AGENT_ID = AgentID("0442ca83-10ce-495f-9c1c-92b4e1f5c39c") @pytest.fixture @@ -25,7 +28,7 @@ def test_uses_correct_extension(ransomware_options_dict, tmp_path, ransomware_fi ransomware_directories["linux_target_dir"] = target_dir ransomware_directories["windows_target_dir"] = target_dir ransomware = ransomware_builder.build_ransomware( - ransomware_options_dict, MagicMock(spec=IAgentEventQueue) + ransomware_options_dict, MagicMock(spec=IAgentEventQueue), AGENT_ID ) file = target_dir / "file.txt" diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 01c32947088..478ce3bc1d9 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -14,6 +14,7 @@ from common.agent_events import AbstractAgentEvent, FileEncryptionEvent from common.event_queue import AgentEventSubscriber, IAgentEventQueue +from common.types import AgentID from infection_monkey.payload.ransomware.consts import README_FILE_NAME, README_SRC from infection_monkey.payload.ransomware.ransomware import Ransomware from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions @@ -64,6 +65,7 @@ def inner( file_selector, leave_readme, agent_event_queue_spy, + AgentID("8f53f4fb-2d33-465a-aa9c-de704a7e42b3"), ) return inner