diff --git a/monkey/common/common_consts/telem_categories.py b/monkey/common/common_consts/telem_categories.py index 280cfce05..dc083d4ab 100644 --- a/monkey/common/common_consts/telem_categories.py +++ b/monkey/common/common_consts/telem_categories.py @@ -8,3 +8,4 @@ class TelemCategoryEnum: TRACE = "trace" TUNNEL = "tunnel" ATTACK = "attack" + RANSOMWARE = "ransomware" diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index abd0b3f18..a70781333 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -25,6 +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.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 @@ -233,7 +236,7 @@ class InfectionMonkey(object): if not self._keep_running: break - RansomewarePayload(WormConfiguration.ransomware).run_payload() + InfectionMonkey.run_ransomware() if (not is_empty) and (WormConfiguration.max_iterations > iteration_index + 1): time_to_sleep = WormConfiguration.timeout_between_iterations @@ -463,3 +466,8 @@ class InfectionMonkey(object): def log_arguments(self): arg_string = " ".join([f"{key}: {value}" for key, value in vars(self._opts).items()]) LOG.info(f"Monkey started with arguments: {arg_string}") + + @staticmethod + def run_ransomware(): + telemetry_messenger = TelemetryMessengerWrapper() + RansomewarePayload(WormConfiguration.ransomware, telemetry_messenger).run_payload() 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/infection_monkey/telemetry/base_telem.py b/monkey/infection_monkey/telemetry/base_telem.py index 0fcf4b203..4a37a9eb9 100644 --- a/monkey/infection_monkey/telemetry/base_telem.py +++ b/monkey/infection_monkey/telemetry/base_telem.py @@ -3,6 +3,7 @@ import json import logging from infection_monkey.control import ControlClient +from infection_monkey.telemetry.i_telem import ITelem logger = logging.getLogger(__name__) LOGGED_DATA_LENGTH = 300 # How many characters of telemetry data will be logged @@ -24,14 +25,11 @@ __author__ = "itay.mizeretz" # logging and sending them. -class BaseTelem(object, metaclass=abc.ABCMeta): +class BaseTelem(ITelem, metaclass=abc.ABCMeta): """ Abstract base class for telemetry. """ - def __init__(self): - pass - def send(self, log_data=True): """ Sends telemetry to island @@ -41,13 +39,6 @@ class BaseTelem(object, metaclass=abc.ABCMeta): self._log_telem_sending(serialized_data, log_data) ControlClient.send_telemetry(self.telem_category, serialized_data) - @abc.abstractmethod - def get_data(self) -> dict: - """ - :return: Data of telemetry (should be dict) - """ - pass - @property def json_encoder(self): return json.JSONEncoder @@ -57,14 +48,6 @@ class BaseTelem(object, metaclass=abc.ABCMeta): if log_data: logger.debug(f"Telemetry contents: {BaseTelem._truncate_data(serialized_data)}") - @property - @abc.abstractmethod - def telem_category(self): - """ - :return: Telemetry type - """ - pass - @staticmethod def _truncate_data(data: str): if len(data) <= LOGGED_DATA_LENGTH: diff --git a/monkey/infection_monkey/telemetry/i_telem.py b/monkey/infection_monkey/telemetry/i_telem.py new file mode 100644 index 000000000..faaa0a65e --- /dev/null +++ b/monkey/infection_monkey/telemetry/i_telem.py @@ -0,0 +1,29 @@ +import abc + + +class ITelem(metaclass=abc.ABCMeta): + @abc.abstractmethod + def send(self, log_data=True): + """ + Sends telemetry to island + """ + + @abc.abstractmethod + def get_data(self) -> dict: + """ + :return: Data of telemetry (should be dict) + """ + pass + + @property + @abc.abstractmethod + def json_encoder(self): + pass + + @property + @abc.abstractmethod + def telem_category(self): + """ + :return: Telemetry type + """ + pass diff --git a/monkey/infection_monkey/telemetry/messengers/i_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/i_telemetry_messenger.py new file mode 100644 index 000000000..cf5511e83 --- /dev/null +++ b/monkey/infection_monkey/telemetry/messengers/i_telemetry_messenger.py @@ -0,0 +1,9 @@ +import abc + +from infection_monkey.telemetry.i_telem import ITelem + + +class ITelemetryMessenger(metaclass=abc.ABCMeta): + @abc.abstractmethod + def send_telemetry(self, telemetry: ITelem): + pass diff --git a/monkey/infection_monkey/telemetry/messengers/telemetry_messenger_wrapper.py b/monkey/infection_monkey/telemetry/messengers/telemetry_messenger_wrapper.py new file mode 100644 index 000000000..e436cdd46 --- /dev/null +++ b/monkey/infection_monkey/telemetry/messengers/telemetry_messenger_wrapper.py @@ -0,0 +1,7 @@ +from infection_monkey.telemetry.i_telem import ITelem +from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger + + +class TelemetryMessengerWrapper(ITelemetryMessenger): + def send_telemetry(self, telemetry: ITelem): + telemetry.send() diff --git a/monkey/infection_monkey/telemetry/ransomware_telem.py b/monkey/infection_monkey/telemetry/ransomware_telem.py new file mode 100644 index 000000000..c56e8337c --- /dev/null +++ b/monkey/infection_monkey/telemetry/ransomware_telem.py @@ -0,0 +1,22 @@ +from typing import List, Tuple + +from common.common_consts.telem_categories import TelemCategoryEnum +from infection_monkey.telemetry.base_telem import BaseTelem + + +class RansomwareTelem(BaseTelem): + def __init__(self, attempts: List[Tuple[str, str]]): + """ + Ransomware telemetry constructor + :param attempts: List of tuples with each tuple containing the path + of a file it tried encrypting and its result. + If ransomware fails completely - list of one tuple + containing the directory path and error string. + """ + super().__init__() + self.attempts = attempts + + telem_category = TelemCategoryEnum.RANSOMWARE + + def get_data(self): + return {"ransomware_attempts": self.attempts} 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] diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py new file mode 100644 index 000000000..4994c9287 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py @@ -0,0 +1,20 @@ +import json + +import pytest + +from infection_monkey.telemetry.ransomware_telem import RansomwareTelem + +ATTEMPTS = [("", "")] + + +@pytest.fixture +def ransomware_telem_test_instance(): + return RansomwareTelem(ATTEMPTS) + + +def test_ransomware_telem_send(ransomware_telem_test_instance, spy_send_telemetry): + ransomware_telem_test_instance.send() + expected_data = {"ransomware_attempts": ATTEMPTS} + expected_data = json.dumps(expected_data, cls=ransomware_telem_test_instance.json_encoder) + assert spy_send_telemetry.data == expected_data + assert spy_send_telemetry.telem_category == "ransomware"