diff --git a/monkey/infection_monkey/ransomware/ransomware_config.py b/monkey/infection_monkey/ransomware/ransomware_config.py new file mode 100644 index 000000000..26805c47c --- /dev/null +++ b/monkey/infection_monkey/ransomware/ransomware_config.py @@ -0,0 +1,25 @@ +import logging + +from common.utils.file_utils import InvalidPath, expand_path +from infection_monkey.utils.environment import is_windows_os + +LOG = logging.getLogger(__name__) + + +class RansomwareConfig: + def __init__(self, config: dict): + self.encryption_enabled = config["encryption"]["enabled"] + self.readme_enabled = config["other_behaviors"]["readme"] + self._set_target_directory(config["encryption"]["directories"]) + + def _set_target_directory(self, os_target_directories: dict): + if is_windows_os(): + target_directory = os_target_directories["windows_target_dir"] + else: + target_directory = os_target_directories["linux_target_dir"] + + try: + self.target_directory = expand_path(target_directory) + except InvalidPath as e: + LOG.debug(f"Target ransomware directory set to None: {e}") + self.target_directory = None diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py new file mode 100644 index 000000000..141186f18 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py @@ -0,0 +1,73 @@ +from pathlib import Path + +import pytest +from tests.utils import raise_ + +from common.utils.file_utils import InvalidPath +from infection_monkey.ransomware import ransomware_config +from infection_monkey.ransomware.ransomware_config import RansomwareConfig + +LINUX_DIR = "/tmp/test" +WINDOWS_DIR = "C:\\tmp\\test" + + +@pytest.fixture +def config_from_island(): + return { + "encryption": { + "enabled": None, + "directories": { + "linux_target_dir": LINUX_DIR, + "windows_target_dir": WINDOWS_DIR, + }, + }, + "other_behaviors": {"readme": None}, + } + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_encryption_enabled(enabled, config_from_island): + config_from_island["encryption"]["enabled"] = enabled + config = RansomwareConfig(config_from_island) + + assert config.encryption_enabled == enabled + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_readme_enabled(enabled, config_from_island): + config_from_island["other_behaviors"]["readme"] = enabled + config = RansomwareConfig(config_from_island) + + assert config.readme_enabled == enabled + + +def test_linux_target_dir(monkeypatch, config_from_island): + monkeypatch.setattr(ransomware_config, "is_windows_os", lambda: False) + + config = RansomwareConfig(config_from_island) + assert config.target_directory == Path(LINUX_DIR) + + +def test_windows_target_dir(monkeypatch, config_from_island): + monkeypatch.setattr(ransomware_config, "is_windows_os", lambda: True) + + config = RansomwareConfig(config_from_island) + assert config.target_directory == Path(WINDOWS_DIR) + + +def test_env_variables_in_target_dir_resolved(config_from_island, patched_home_env, tmp_path): + path_with_env_variable = "$HOME/ransomware_target" + + config_from_island["encryption"]["directories"]["linux_target_dir"] = config_from_island[ + "encryption" + ]["directories"]["windows_target_dir"] = path_with_env_variable + + config = RansomwareConfig(config_from_island) + assert config.target_directory == patched_home_env / "ransomware_target" + + +def test_target_dir_is_none(monkeypatch, config_from_island): + monkeypatch.setattr(ransomware_config, "expand_path", lambda _: raise_(InvalidPath("invalid"))) + + config = RansomwareConfig(config_from_island) + assert config.target_directory is None