p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Fix custom PBA related unit tests

This commit is contained in:
vakarisz 2022-03-30 15:48:43 +03:00
parent 1f31e96adb
commit 2e3a718469
2 changed files with 81 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

149
monkey/tests/unit_tests/infection_monkey/post_breach/actions/test_users_custom_pba.py
View File

@ -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 == ""

19
monkey/tests/unit_tests/monkey_island/cc/services/test_config.py
View File

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