diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_options.py b/monkey/infection_monkey/payload/ransomware/ransomware_options.py index 8416f8465..505974ae0 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_options.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_options.py @@ -9,6 +9,7 @@ logger = logging.getLogger(__name__) class RansomwareOptions: def __init__(self, options: dict): self.encryption_enabled = options["encryption"]["enabled"] + self.file_extension = options["encryption"]["file_extension"] self.readme_enabled = options["other_behaviors"]["readme"] self.target_directory = None diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py index f2b6a8c8c..ae7280a86 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_options.py @@ -7,6 +7,7 @@ from common.utils.file_utils import InvalidPath from infection_monkey.payload.ransomware import ransomware_options from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions +EXTENSION = ".testext" LINUX_DIR = "/tmp/test" WINDOWS_DIR = "C:\\tmp\\test" @@ -16,6 +17,7 @@ def options_from_island(): return { "encryption": { "enabled": None, + "file_extension": EXTENSION, "directories": { "linux_target_dir": LINUX_DIR, "windows_target_dir": WINDOWS_DIR, @@ -41,6 +43,12 @@ def test_readme_enabled(enabled, options_from_island): assert options.readme_enabled == enabled +def test_file_extension(options_from_island): + options = RansomwareOptions(options_from_island) + + assert options.file_extension == EXTENSION + + def test_linux_target_dir(monkeypatch, options_from_island): monkeypatch.setattr(ransomware_options, "is_windows_os", lambda: False)