diff --git a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py index e8cdd0333..598f3412b 100644 --- a/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py +++ b/monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py @@ -2,19 +2,20 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.post_breach.custom_pba.users_custom_pba import UsersPBA +from infection_monkey.post_breach.custom_pba.custom_pba import CustomPBA MONKEY_DIR_PATH = "/dir/to/monkey/" CUSTOM_LINUX_CMD = "command-for-linux" CUSTOM_LINUX_FILENAME = "filename-for-linux" CUSTOM_WINDOWS_CMD = "command-for-windows" CUSTOM_WINDOWS_FILENAME = "filename-for-windows" +CUSTOM_SERVER = "10.10.10.10:5000" -@pytest.fixture +@pytest.fixture(autouse=True) def fake_monkey_dir_path(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.get_monkey_dir_path", + "infection_monkey.post_breach.custom_pba.custom_pba.get_monkey_dir_path", lambda: MONKEY_DIR_PATH, ) @@ -22,7 +23,7 @@ def fake_monkey_dir_path(monkeypatch): @pytest.fixture def set_os_linux(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os", lambda: False, ) @@ -30,106 +31,92 @@ def set_os_linux(monkeypatch): @pytest.fixture def set_os_windows(monkeypatch): monkeypatch.setattr( - "infection_monkey.post_breach.actions.users_custom_pba.is_windows_os", + "infection_monkey.post_breach.custom_pba.custom_pba.is_windows_os", lambda: True, ) @pytest.fixture -def mock_UsersPBA_linux_custom_file_and_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", - CUSTOM_LINUX_CMD, - ) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_linux_filename", - CUSTOM_LINUX_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_custom_pba_linux_options(): + return { + "linux_command": CUSTOM_LINUX_CMD, + "linux_filename": CUSTOM_LINUX_FILENAME, + "windows_command": "", + "windows_filename": "", + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_linux_custom_file_and_cmd( - mock_UsersPBA_linux_custom_file_and_cmd, -): +def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_custom_pba_linux_options) expected_command = f"cd {MONKEY_DIR_PATH} ; {CUSTOM_LINUX_CMD}" - assert mock_UsersPBA_linux_custom_file_and_cmd.command == expected_command + assert pba.command == expected_command + assert pba.filename == CUSTOM_LINUX_FILENAME @pytest.fixture -def mock_UsersPBA_windows_custom_file_and_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", - CUSTOM_WINDOWS_CMD, - ) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_windows_filename", - CUSTOM_WINDOWS_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_custom_pba_windows_options(): + return { + "linux_command": "", + "linux_filename": "", + "windows_command": CUSTOM_WINDOWS_CMD, + "windows_filename": CUSTOM_WINDOWS_FILENAME, + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_windows_custom_file_and_cmd( - mock_UsersPBA_windows_custom_file_and_cmd, -): +def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows): + + pba = CustomPBA(MagicMock()) + pba._set_options(fake_custom_pba_windows_options) expected_command = f"cd {MONKEY_DIR_PATH} & {CUSTOM_WINDOWS_CMD}" - assert mock_UsersPBA_windows_custom_file_and_cmd.command == expected_command + assert pba.command == expected_command + assert pba.filename == CUSTOM_WINDOWS_FILENAME @pytest.fixture -def mock_UsersPBA_linux_custom_file(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_linux_filename", - CUSTOM_LINUX_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_options_files_only(): + return { + "linux_command": "", + "linux_filename": CUSTOM_LINUX_FILENAME, + "windows_command": "", + "windows_filename": CUSTOM_WINDOWS_FILENAME, + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file): - expected_command = "" - assert mock_UsersPBA_linux_custom_file.command == expected_command +@pytest.mark.parametrize("os", [set_os_linux, set_os_windows]) +def test_files_only(fake_options_files_only, os): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_files_only) + assert pba.command == "" @pytest.fixture -def mock_UsersPBA_windows_custom_file(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None) - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.PBA_windows_filename", - CUSTOM_WINDOWS_FILENAME, - ) - return UsersPBA(MagicMock()) +def fake_options_commands_only(): + return { + "linux_command": CUSTOM_LINUX_CMD, + "linux_filename": "", + "windows_command": CUSTOM_WINDOWS_CMD, + "windows_filename": "", + # Current server is used for attack telemetry + "current_server": CUSTOM_SERVER, + } -def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file): - expected_command = "" - assert mock_UsersPBA_windows_custom_file.command == expected_command +def test_commands_only(fake_options_commands_only, set_os_linux): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_commands_only) + assert pba.command == CUSTOM_LINUX_CMD + assert pba.filename == "" -@pytest.fixture -def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", - CUSTOM_LINUX_CMD, - ) - monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_linux_filename", None) - return UsersPBA(MagicMock()) - - -def test_command_linux_custom_cmd(mock_UsersPBA_linux_custom_cmd): - expected_command = CUSTOM_LINUX_CMD - assert mock_UsersPBA_linux_custom_cmd.command == expected_command - - -@pytest.fixture -def mock_UsersPBA_windows_custom_cmd(set_os_windows, fake_monkey_dir_path, monkeypatch): - monkeypatch.setattr( - "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", - CUSTOM_WINDOWS_CMD, - ) - monkeypatch.setattr("infection_monkey.config.WormConfiguration.PBA_windows_filename", None) - return UsersPBA(MagicMock()) - - -def test_command_windows_custom_cmd(mock_UsersPBA_windows_custom_cmd): - expected_command = CUSTOM_WINDOWS_CMD - assert mock_UsersPBA_windows_custom_cmd.command == expected_command +def test_commands_only_windows(fake_options_commands_only, set_os_windows): + pba = CustomPBA(MagicMock()) + pba._set_options(fake_options_commands_only) + assert pba.command == CUSTOM_WINDOWS_CMD + assert pba.filename == "" diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py index ae0a44cdc..b49007eb0 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_config.py @@ -66,12 +66,6 @@ def test_format_config_for_agent__pbas(flat_monkey_config): "ScheduleJobs": {}, "Timestomping": {}, "AccountDiscovery": {}, - "Custom": { - "linux_command": "bash test.sh", - "windows_command": "powershell test.ps1", - "linux_filename": "test.sh", - "windows_filename": "test.ps1", - }, } ConfigService.format_flat_config_for_agent(flat_monkey_config) @@ -84,6 +78,19 @@ def test_format_config_for_agent__pbas(flat_monkey_config): assert "PBA_windows_filename" not in flat_monkey_config +def test_format_config_for_custom_pbas(flat_monkey_config): + custom_config = { + "linux_command": "bash test.sh", + "windows_command": "powershell test.ps1", + "linux_filename": "test.sh", + "windows_filename": "test.ps1", + "current_server": "10.197.94.72:5000", + } + ConfigService.format_flat_config_for_agent(flat_monkey_config) + + assert flat_monkey_config["custom_pbas"] == custom_config + + def test_get_config_propagation_credentials_from_flat_config(flat_monkey_config): expected_creds = { "exploit_lm_hash_list": ["lm_hash_1", "lm_hash_2"],