Agent: Pass Agent ID to Ransomware

Issue #3119 PR #??
guardicore · Mar 24, 2023 · 4aed973 · 4aed973
1 parent 59fed74
commit 4aed973
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 8 deletions.
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
@@ -459,7 +459,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
@@ -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,13 +26,15 @@ def __init__(
         select_files: Callable[[Path], Iterable[Path]],
         leave_readme: Callable[[Path, Path], None],
         agent_event_queue: IAgentEventQueue,
+        agent_id: AgentID,
     ):
         self._config = config
 
         self._encrypt_file = encrypt_file
         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
@@ -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
@@ -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/...tests/integration_tests/infection_monkey/payload/ransomware/test_integrated_ransomware.py b/...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
@@ -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