From 6f5a7faaa194b1eb0ed1f3c7e07dea885b35b2b0 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 15 Jul 2021 10:05:34 -0400
Subject: [PATCH] Agent: Add RannsomwareConfig class

---
 .../ransomware/ransomware_config.py           | 25 +++++++
 .../ransomware/test_ransomware_config.py      | 73 +++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 monkey/infection_monkey/ransomware/ransomware_config.py
 create mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_config.py

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