Agent: Fix custom PBA related unit tests

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):
+def fake_custom_pba_linux_options():
-    monkeypatch.setattr(
+    return {
-        "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
+        "linux_command": CUSTOM_LINUX_CMD,
-        CUSTOM_LINUX_CMD,
+        "linux_filename": CUSTOM_LINUX_FILENAME,
-    )
+        "windows_command": "",
-    monkeypatch.setattr(
+        "windows_filename": "",
-        "infection_monkey.config.WormConfiguration.PBA_linux_filename",
+        # Current server is used for attack telemetry
-        CUSTOM_LINUX_FILENAME,
+        "current_server": CUSTOM_SERVER,
-    )
+    }
-    return UsersPBA(MagicMock())
 
 
-def test_command_linux_custom_file_and_cmd(
+def test_command_linux_custom_file_and_cmd(fake_custom_pba_linux_options, set_os_linux):
-    mock_UsersPBA_linux_custom_file_and_cmd,
+    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):
+def fake_custom_pba_windows_options():
-    monkeypatch.setattr(
+    return {
-        "infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd",
+        "linux_command": "",
-        CUSTOM_WINDOWS_CMD,
+        "linux_filename": "",
-    )
+        "windows_command": CUSTOM_WINDOWS_CMD,
-    monkeypatch.setattr(
+        "windows_filename": CUSTOM_WINDOWS_FILENAME,
-        "infection_monkey.config.WormConfiguration.PBA_windows_filename",
+        # Current server is used for attack telemetry
-        CUSTOM_WINDOWS_FILENAME,
+        "current_server": CUSTOM_SERVER,
-    )
+    }
-    return UsersPBA(MagicMock())
 
 
-def test_command_windows_custom_file_and_cmd(
+def test_command_windows_custom_file_and_cmd(fake_custom_pba_windows_options, set_os_windows):
-    mock_UsersPBA_windows_custom_file_and_cmd,
+
-):
+    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):
+def fake_options_files_only():
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd", None)
+    return {
-    monkeypatch.setattr(
+        "linux_command": "",
-        "infection_monkey.config.WormConfiguration.PBA_linux_filename",
+        "linux_filename": CUSTOM_LINUX_FILENAME,
-        CUSTOM_LINUX_FILENAME,
+        "windows_command": "",
-    )
+        "windows_filename": CUSTOM_WINDOWS_FILENAME,
-    return UsersPBA(MagicMock())
+        # Current server is used for attack telemetry
+        "current_server": CUSTOM_SERVER,
+    }
 
 
-def test_command_linux_custom_file(mock_UsersPBA_linux_custom_file):
+@pytest.mark.parametrize("os", [set_os_linux, set_os_windows])
-    expected_command = ""
+def test_files_only(fake_options_files_only, os):
-    assert mock_UsersPBA_linux_custom_file.command == expected_command
+    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):
+def fake_options_commands_only():
-    monkeypatch.setattr("infection_monkey.config.WormConfiguration.custom_PBA_windows_cmd", None)
+    return {
-    monkeypatch.setattr(
+        "linux_command": CUSTOM_LINUX_CMD,
-        "infection_monkey.config.WormConfiguration.PBA_windows_filename",
+        "linux_filename": "",
-        CUSTOM_WINDOWS_FILENAME,
+        "windows_command": CUSTOM_WINDOWS_CMD,
-    )
+        "windows_filename": "",
-    return UsersPBA(MagicMock())
+        # Current server is used for attack telemetry
+        "current_server": CUSTOM_SERVER,
+    }
 
 
-def test_command_windows_custom_file(mock_UsersPBA_windows_custom_file):
+def test_commands_only(fake_options_commands_only, set_os_linux):
-    expected_command = ""
+    pba = CustomPBA(MagicMock())
-    assert mock_UsersPBA_windows_custom_file.command == expected_command
+    pba._set_options(fake_options_commands_only)
+    assert pba.command == CUSTOM_LINUX_CMD
+    assert pba.filename == ""
 
 
-@pytest.fixture
+def test_commands_only_windows(fake_options_commands_only, set_os_windows):
-def mock_UsersPBA_linux_custom_cmd(set_os_linux, fake_monkey_dir_path, monkeypatch):
+    pba = CustomPBA(MagicMock())
-    monkeypatch.setattr(
+    pba._set_options(fake_options_commands_only)
-        "infection_monkey.config.WormConfiguration.custom_PBA_linux_cmd",
+    assert pba.command == CUSTOM_WINDOWS_CMD
-        CUSTOM_LINUX_CMD,
+    assert pba.filename == ""
-    )
-    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

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"],