agent: Send telemetry from ransomware payload

monkey/infection_monkey/monkey.py

@@ -25,7 +25,9 @@ from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
 from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem
-from infection_monkey.telemetry.ransomware_telem import RansomwareTelem
+from infection_monkey.telemetry.messengers.telemetry_messenger_wrapper import (
+    TelemetryMessengerWrapper,
+)
 from infection_monkey.telemetry.scan_telem import ScanTelem
 from infection_monkey.telemetry.state_telem import StateTelem
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
@@ -234,8 +236,8 @@ class InfectionMonkey(object):
                     if not self._keep_running:
                         break
 
-                ransomware_attempts = RansomewarePayload(WormConfiguration.ransomware).run_payload()
-                RansomwareTelem(ransomware_attempts).send()
+                telemetry_messenger = TelemetryMessengerWrapper()
+                RansomewarePayload(WormConfiguration.ransomware, telemetry_messenger).run_payload()
 
                 if (not is_empty) and (WormConfiguration.max_iterations > iteration_index + 1):
                     time_to_sleep = WormConfiguration.timeout_between_iterations

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -5,6 +5,8 @@ from typing import List, Optional, Tuple
 from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor
 from infection_monkey.ransomware.file_selectors import select_production_safe_target_files
 from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
+from infection_monkey.telemetry.ransomware_telem import RansomwareTelem
 from infection_monkey.utils.environment import is_windows_os
 
 LOG = logging.getLogger(__name__)
@@ -14,7 +16,7 @@ CHUNK_SIZE = 4096 * 24
 
 
 class RansomewarePayload:
-    def __init__(self, config: dict):
+    def __init__(self, config: dict, telemetry_messenger: ITelemetryMessenger):
         LOG.info(f"Windows dir configured for encryption is \"{config['windows_dir']}\"")
         LOG.info(f"Linux dir configured for encryption is \"{config['linux_dir']}\"")
 
@@ -25,6 +27,7 @@ class RansomewarePayload:
         self._valid_file_extensions_for_encryption.discard(self._new_file_extension)
 
         self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE)
+        self._telemetry_messenger = telemetry_messenger
 
     def run_payload(self):
         file_list = self._find_files()
@@ -44,12 +47,16 @@ class RansomewarePayload:
             try:
                 self._encryptor.encrypt_file_in_place(filepath)
                 self._add_extension(filepath)
-                results.append((filepath, None))
+                self._send_telemetry(filepath, "")
             except Exception as ex:
-                results.append((filepath, ex))
+                self._send_telemetry(filepath, str(ex))
 
         return results
 
     def _add_extension(self, filepath: Path):
         new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
         filepath.rename(new_filepath)
+
+    def _send_telemetry(self, filepath: Path, error: str):
+        encryption_attempt = RansomwareTelem((str(filepath), str(error)))
+        self._telemetry_messenger.send_telemetry(encryption_attempt)

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import PurePath
 
 import pytest
 from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
@@ -20,7 +21,18 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
 )
 from tests.utils import hash_file, is_user_admin
 
+from infection_monkey.ransomware import ransomware_payload as ransomware_payload_module
 from infection_monkey.ransomware.ransomware_payload import EXTENSION, RansomewarePayload
+from infection_monkey.telemetry.i_telem import ITelem
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
+
+
+class TelemetryMessengerSpy(ITelemetryMessenger):
+    def __init__(self):
+        self.telemetries = []
+
+    def send_telemetry(self, telemetry: ITelem):
+        self.telemetries.append(telemetry)
 
 
 def with_extension(filename):
@@ -33,8 +45,13 @@ def ransomware_payload_config(ransomware_target):
 
 
 @pytest.fixture
-def ransomware_payload(ransomware_payload_config):
-    return RansomewarePayload(ransomware_payload_config)
+def telemetry_messenger_spy():
+    return TelemetryMessengerSpy()
+
+
+@pytest.fixture
+def ransomware_payload(ransomware_payload_config, telemetry_messenger_spy):
+    return RansomewarePayload(ransomware_payload_config, telemetry_messenger_spy)
 
 
 def test_file_with_excluded_extension_not_encrypted(ransomware_target, ransomware_payload):
@@ -120,3 +137,30 @@ def test_skip_already_encrypted_file(ransomware_target, ransomware_payload):
         hash_file(ransomware_target / ALREADY_ENCRYPTED_TXT_M0NK3Y)
         == ALREADY_ENCRYPTED_TXT_M0NK3Y_CLEARTEXT_SHA256
     )
+
+
+def test_telemetry_success(ransomware_payload, telemetry_messenger_spy):
+    ransomware_payload.run_payload()
+
+    assert len(telemetry_messenger_spy.telemetries) == 2
+    telem_1 = telemetry_messenger_spy.telemetries[0]
+    telem_2 = telemetry_messenger_spy.telemetries[1]
+
+    assert ALL_ZEROS_PDF in telem_1.get_data()["ransomware_attempts"][0]
+    assert telem_1.get_data()["ransomware_attempts"][1] == ""
+    assert TEST_KEYBOARD_TXT in telem_2.get_data()["ransomware_attempts"][0]
+    assert telem_2.get_data()["ransomware_attempts"][1] == ""
+
+
+def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_spy):
+    monkeypatch.setattr(
+        ransomware_payload_module,
+        "select_production_safe_target_files",
+        lambda a, b: [PurePath("/file/not/exist")],
+    ),
+
+    ransomware_payload.run_payload()
+    telem_1 = telemetry_messenger_spy.telemetries[0]
+
+    assert "/file/not/exist" in telem_1.get_data()["ransomware_attempts"][0]
+    assert "No such file or directory" in telem_1.get_data()["ransomware_attempts"][1]