Agent, UT: Rename 'config' to 'options' in ransomware files

This commit is contained in:
Shreya Malviya 2021-12-17 19:55:26 +05:30
parent 0328d2860e
commit 958cf3a252
7 changed files with 117 additions and 117 deletions

View File

@ -7,7 +7,7 @@ from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
from .consts import README_FILE_NAME, README_SRC
from .ransomware_config import RansomwareConfig
from .ransomware_options import RansomwareOptions
logger = logging.getLogger(__name__)
@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
class Ransomware:
def __init__(
self,
config: RansomwareConfig,
config: RansomwareOptions,
encrypt_file: Callable[[Path], None],
select_files: Callable[[Path], List[Path]],
leave_readme: Callable[[Path, Path], None],

View File

@ -13,7 +13,7 @@ from . import readme_dropper
from .file_selectors import ProductionSafeTargetFileSelector
from .in_place_file_encryptor import InPlaceFileEncryptor
from .ransomware import Ransomware
from .ransomware_config import RansomwareConfig
from .ransomware_options import RansomwareOptions
from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS
EXTENSION = ".m0nk3y"
@ -22,9 +22,9 @@ CHUNK_SIZE = 4096 * 24
logger = logging.getLogger(__name__)
def build_ransomware(config: dict):
logger.debug(f"Ransomware configuration:\n{pformat(config)}")
ransomware_config = RansomwareConfig(config)
def build_ransomware(options: dict):
logger.debug(f"Ransomware configuration:\n{pformat(options)}")
ransomware_options = RansomwareOptions(options)
file_encryptor = _build_file_encryptor()
file_selector = _build_file_selector()
@ -32,7 +32,7 @@ def build_ransomware(config: dict):
telemetry_messenger = _build_telemetry_messenger()
return Ransomware(
ransomware_config,
ransomware_options,
file_encryptor,
file_selector,
leave_readme,

View File

@ -6,13 +6,13 @@ from infection_monkey.utils.environment import is_windows_os
logger = logging.getLogger(__name__)
class RansomwareConfig:
def __init__(self, config: dict):
self.encryption_enabled = config["encryption"]["enabled"]
self.readme_enabled = config["other_behaviors"]["readme"]
class RansomwareOptions:
def __init__(self, options: dict):
self.encryption_enabled = options["encryption"]["enabled"]
self.readme_enabled = options["other_behaviors"]["readme"]
self.target_directory = None
self._set_target_directory(config["encryption"]["directories"])
self._set_target_directory(options["encryption"]["directories"])
def _set_target_directory(self, os_target_directories: dict):
if is_windows_os():

View File

@ -10,12 +10,12 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file
from infection_monkey.payload.ransomware.consts import README_FILE_NAME, README_SRC
from infection_monkey.payload.ransomware.ransomware import Ransomware
from infection_monkey.payload.ransomware.ransomware_config import RansomwareConfig
from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions
@pytest.fixture
def ransomware(build_ransomware, ransomware_config):
return build_ransomware(ransomware_config)
def ransomware(build_ransomware, ransomware_options):
return build_ransomware(ransomware_options)
@pytest.fixture
@ -40,14 +40,14 @@ def build_ransomware(
@pytest.fixture
def ransomware_config(ransomware_test_data):
class RansomwareConfigStub(RansomwareConfig):
def ransomware_options(ransomware_test_data):
class RansomwareOptionsStub(RansomwareOptions):
def __init__(self, encryption_enabled, readme_enabled, target_directory):
self.encryption_enabled = encryption_enabled
self.readme_enabled = readme_enabled
self.target_directory = target_directory
return RansomwareConfigStub(True, False, ransomware_test_data)
return RansomwareOptionsStub(True, False, ransomware_test_data)
@pytest.fixture
@ -71,11 +71,11 @@ def mock_leave_readme():
def test_files_selected_from_target_dir(
ransomware,
ransomware_config,
ransomware_options,
mock_file_selector,
):
ransomware.run(threading.Event())
mock_file_selector.assert_called_with(ransomware_config.target_directory)
mock_file_selector.assert_called_with(ransomware_options.target_directory)
def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor):
@ -87,23 +87,23 @@ def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_fil
def test_encryption_skipped_if_configured_false(
build_ransomware, ransomware_config, mock_file_encryptor
build_ransomware, ransomware_options, mock_file_encryptor
):
ransomware_config.encryption_enabled = False
ransomware_options.encryption_enabled = False
ransomware = build_ransomware(ransomware_config)
ransomware = build_ransomware(ransomware_options)
ransomware.run(threading.Event())
assert mock_file_encryptor.call_count == 0
def test_encryption_skipped_if_no_directory(
build_ransomware, ransomware_config, mock_file_encryptor
build_ransomware, ransomware_options, mock_file_encryptor
):
ransomware_config.encryption_enabled = True
ransomware_config.target_directory = None
ransomware_options.encryption_enabled = True
ransomware_options.target_directory = None
ransomware = build_ransomware(ransomware_config)
ransomware = build_ransomware(ransomware_options)
ransomware.run(threading.Event())
assert mock_file_encryptor.call_count == 0
@ -124,13 +124,13 @@ def test_telemetry_success(ransomware, telemetry_messenger_spy):
assert telem_2.get_data()["files"][0]["error"] == ""
def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messenger_spy):
def test_telemetry_failure(build_ransomware, ransomware_options, telemetry_messenger_spy):
file_not_exists = "/file/not/exist"
mfe = MagicMock(
side_effect=FileNotFoundError(f"[Errno 2] No such file or directory: '{file_not_exists}'")
)
mfs = MagicMock(return_value=[PurePosixPath(file_not_exists)])
ransomware = build_ransomware(config=ransomware_config, file_encryptor=mfe, file_selector=mfs)
ransomware = build_ransomware(config=ransomware_options, file_encryptor=mfe, file_selector=mfs)
ransomware.run(threading.Event())
telem = telemetry_messenger_spy.telemetries[0]
@ -140,36 +140,36 @@ def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messen
assert "No such file or directory" in telem.get_data()["files"][0]["error"]
def test_readme_false(build_ransomware, ransomware_config, mock_leave_readme):
ransomware_config.readme_enabled = False
ransomware = build_ransomware(ransomware_config)
def test_readme_false(build_ransomware, ransomware_options, mock_leave_readme):
ransomware_options.readme_enabled = False
ransomware = build_ransomware(ransomware_options)
ransomware.run(threading.Event())
mock_leave_readme.assert_not_called()
def test_readme_true(build_ransomware, ransomware_config, mock_leave_readme, ransomware_test_data):
ransomware_config.readme_enabled = True
ransomware = build_ransomware(ransomware_config)
def test_readme_true(build_ransomware, ransomware_options, mock_leave_readme, ransomware_test_data):
ransomware_options.readme_enabled = True
ransomware = build_ransomware(ransomware_options)
ransomware.run(threading.Event())
mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME)
def test_no_readme_if_no_directory(build_ransomware, ransomware_config, mock_leave_readme):
ransomware_config.target_directory = None
ransomware_config.readme_enabled = True
def test_no_readme_if_no_directory(build_ransomware, ransomware_options, mock_leave_readme):
ransomware_options.target_directory = None
ransomware_options.readme_enabled = True
ransomware = build_ransomware(ransomware_config)
ransomware = build_ransomware(ransomware_options)
ransomware.run(threading.Event())
mock_leave_readme.assert_not_called()
def test_leave_readme_exceptions_handled(build_ransomware, ransomware_config):
def test_leave_readme_exceptions_handled(build_ransomware, ransomware_options):
leave_readme = MagicMock(side_effect=Exception("Test exception when leaving README"))
ransomware_config.readme_enabled = True
ransomware = build_ransomware(config=ransomware_config, leave_readme=leave_readme)
ransomware_options.readme_enabled = True
ransomware = build_ransomware(config=ransomware_options, leave_readme=leave_readme)
# Test will fail if exception is raised and not handled
ransomware.run(threading.Event())

View File

@ -1,73 +0,0 @@
from pathlib import Path
import pytest
from tests.utils import raise_
from common.utils.file_utils import InvalidPath
from infection_monkey.payload.ransomware import ransomware_config
from infection_monkey.payload.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

View File

@ -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.payload.ransomware import ransomware_options
from infection_monkey.payload.ransomware.ransomware_options import RansomwareOptions
LINUX_DIR = "/tmp/test"
WINDOWS_DIR = "C:\\tmp\\test"
@pytest.fixture
def options_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, options_from_island):
options_from_island["encryption"]["enabled"] = enabled
options = RansomwareOptions(options_from_island)
assert options.encryption_enabled == enabled
@pytest.mark.parametrize("enabled", [True, False])
def test_readme_enabled(enabled, options_from_island):
options_from_island["other_behaviors"]["readme"] = enabled
options = RansomwareOptions(options_from_island)
assert options.readme_enabled == enabled
def test_linux_target_dir(monkeypatch, options_from_island):
monkeypatch.setattr(ransomware_options, "is_windows_os", lambda: False)
options = RansomwareOptions(options_from_island)
assert options.target_directory == Path(LINUX_DIR)
def test_windows_target_dir(monkeypatch, options_from_island):
monkeypatch.setattr(ransomware_options, "is_windows_os", lambda: True)
options = RansomwareOptions(options_from_island)
assert options.target_directory == Path(WINDOWS_DIR)
def test_env_variables_in_target_dir_resolved(options_from_island, patched_home_env, tmp_path):
path_with_env_variable = "$HOME/ransomware_target"
options_from_island["encryption"]["directories"]["linux_target_dir"] = options_from_island[
"encryption"
]["directories"]["windows_target_dir"] = path_with_env_variable
options = RansomwareOptions(options_from_island)
assert options.target_directory == patched_home_env / "ransomware_target"
def test_target_dir_is_none(monkeypatch, options_from_island):
monkeypatch.setattr(ransomware_options, "expand_path", lambda _: raise_(InvalidPath("invalid")))
options = RansomwareOptions(options_from_island)
assert options.target_directory is None

View File

@ -38,7 +38,7 @@ def test_format_config_for_agent__credentials_removed(flat_monkey_config):
def test_format_config_for_agent__ransomware_payload(flat_monkey_config):
expected_ransomware_config = {
expected_ransomware_options = {
"ransomware": {
"encryption": {
"enabled": True,
@ -54,7 +54,7 @@ def test_format_config_for_agent__ransomware_payload(flat_monkey_config):
ConfigService.format_flat_config_for_agent(flat_monkey_config)
assert "payloads" in flat_monkey_config
assert flat_monkey_config["payloads"] == expected_ransomware_config
assert flat_monkey_config["payloads"] == expected_ransomware_options
assert "ransomware" not in flat_monkey_config