forked from p34709852/monkey
Agent, UT: Rename 'config' to 'options' in ransomware files
This commit is contained in:
parent
0328d2860e
commit
958cf3a252
|
@ -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],
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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():
|
|
@ -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())
|
||||
|
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in New Issue