Merge pull request #1259 from guardicore/ransomware-telemetry

Ransomware telemetry

monkey/common/common_consts/telem_categories.py

@@ -8,3 +8,4 @@ class TelemCategoryEnum:
     TRACE = "trace"
     TUNNEL = "tunnel"
     ATTACK = "attack"
+    RANSOMWARE = "ransomware"

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()

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/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:

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

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

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()

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}

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]

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 = [("<file>", "<encryption attempt result>")]
+
+
+@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"