diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 0b45d3fbd..122e9e459 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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 diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 460b0fb4c..49525902b 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/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) diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py index d5a155f48..138c60004 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py +++ b/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]