From ca1712cdd65e86d2ed4b8221d4b369e4a19dbde2 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 7 Jul 2021 11:40:56 +0300 Subject: [PATCH] Extract the logic of determining target directory for ransomware payload into a separate method --- .../ransomware/ransomware_payload.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py index a58b50545..0cf8d77ec 100644 --- a/monkey/infection_monkey/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/ransomware/ransomware_payload.py @@ -4,7 +4,7 @@ from pathlib import Path from pprint import pformat from typing import List, Optional, Tuple -from common.utils.file_utils import expand_path +from common.utils.file_utils import InvalidPath, expand_path from infection_monkey.ransomware.bitflip_encryptor import BitflipEncryptor from infection_monkey.ransomware.file_selectors import select_production_safe_target_files from infection_monkey.ransomware.targeted_file_extensions import TARGETED_FILE_EXTENSIONS @@ -28,15 +28,7 @@ class RansomwarePayload: self._encryption_enabled = config["encryption"]["enabled"] self._readme_enabled = config["other_behaviors"]["readme"] - target_directories = config["encryption"]["directories"] - self._target_dir = Path( - expand_path( - target_directories["windows_target_dir"] - if is_windows_os() - else target_directories["linux_target_dir"] - ) - ) - + self._target_dir = RansomwarePayload.get_target_dir(config) self._new_file_extension = EXTENSION self._valid_file_extensions_for_encryption = TARGETED_FILE_EXTENSIONS.copy() self._valid_file_extensions_for_encryption.discard(self._new_file_extension) @@ -44,6 +36,19 @@ class RansomwarePayload: self._encryptor = BitflipEncryptor(chunk_size=CHUNK_SIZE) self._telemetry_messenger = telemetry_messenger + @staticmethod + def get_target_dir(config: dict): + target_directories = config["encryption"]["directories"] + if is_windows_os(): + target_dir_field = target_directories["windows_target_dir"] + else: + target_dir_field = target_directories["linux_target_dir"] + + try: + return Path(expand_path(target_dir_field)) + except InvalidPath: + return None + def run_payload(self): if self._encryption_enabled: LOG.info("Running ransomware payload")