diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py index 73ab7b7bc..742f33f56 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -1,4 +1,3 @@ -from collections import namedtuple from unittest.mock import MagicMock import pytest @@ -15,57 +14,57 @@ USER_LIST = ["user1", "user2"] PASSWORD_LIST = ["pass1", "pass2"] LM_HASH_LIST = ["bogo_lm_1"] NT_HASH_LIST = ["bogo_nt_1", "bogo_nt_2"] -DROPPER_TARGET_PATH_32 = "C:\\agent32" DROPPER_TARGET_PATH_64 = "C:\\agent64" -Config = namedtuple( - "Config", - [ - "exploit_user_list", - "exploit_password_list", - "exploit_lm_hash_list", - "exploit_ntlm_hash_list", - "dropper_target_path_win_64", - ], -) - class AuthenticationErrorForTests(Exception): pass +@pytest.fixture +def powershell_arguments(): + options = { + "dropper_target_path_win_64": DROPPER_TARGET_PATH_64, + "credentials": { + "exploit_user_list": USER_LIST, + "exploit_password_list": PASSWORD_LIST, + "exploit_lm_hash_list": LM_HASH_LIST, + "exploit_ntlm_hash_list": NT_HASH_LIST, + }, + } + arguments = { + "host": VictimHost("127.0.0.1"), + "options": options, + "current_depth": 2, + "telemetry_messenger": MagicMock(), + "agent_repository": MagicMock(), + } + return arguments + + @pytest.fixture def powershell_exploiter(monkeypatch): - host = VictimHost("127.0.0.1") - pe = powershell.PowerShellExploiter(host) - pe._config = Config( - USER_LIST, - PASSWORD_LIST, - LM_HASH_LIST, - NT_HASH_LIST, - DROPPER_TARGET_PATH_32, - DROPPER_TARGET_PATH_64, - ) + pe = powershell.PowerShellExploiter() monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests) monkeypatch.setattr(powershell, "is_windows_os", lambda: True) # It's regrettable to mock out a private method on the PowerShellExploiter instance object, but # it's necessary to avoid having to deal with the monkeyfs. TODO: monkeyfs has been removed, so # fix this. - monkeypatch.setattr(pe, "_write_virtual_file_to_local_path", lambda: None) + # monkeypatch.setattr(pe, "_create_local_agent_file", lambda: None) return pe -def test_powershell_disabled(monkeypatch, powershell_exploiter): +def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_arguments): mock_powershell_client = MagicMock(side_effect=Exception) monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) - success = powershell_exploiter.exploit_host() - assert not success + exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) + assert not exploit_result.exploitation_success -def test_powershell_http(monkeypatch, powershell_exploiter): +def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments): def allow_http(_, credentials: Credentials, auth_options: AuthOptions): if not auth_options.ssl: raise AuthenticationErrorForTests @@ -74,13 +73,13 @@ def test_powershell_http(monkeypatch, powershell_exploiter): mock_powershell_client = MagicMock(side_effect=allow_http) monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) - powershell_exploiter.exploit_host() + powershell_exploiter.exploit_host(**powershell_arguments) for call_args in mock_powershell_client.call_args_list: assert not call_args[0][2].ssl -def test_powershell_https(monkeypatch, powershell_exploiter): +def test_powershell_https(monkeypatch, powershell_exploiter, powershell_arguments): def allow_https(_, credentials: Credentials, auth_options: AuthOptions): if auth_options.ssl: raise AuthenticationErrorForTests @@ -89,19 +88,19 @@ def test_powershell_https(monkeypatch, powershell_exploiter): mock_powershell_client = MagicMock(side_effect=allow_https) monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) - powershell_exploiter.exploit_host() + powershell_exploiter.exploit_host(**powershell_arguments) for call_args in mock_powershell_client.call_args_list: if call_args[0][1].secret != "" and call_args[0][1].secret != "dummy_password": assert call_args[0][2].ssl -def test_no_valid_credentials(monkeypatch, powershell_exploiter): +def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_arguments): mock_powershell_client = MagicMock(side_effect=AuthenticationErrorForTests) monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) - success = powershell_exploiter.exploit_host() - assert not success + exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) + assert not exploit_result.exploitation_success def authenticate(mock_client): @@ -114,29 +113,29 @@ def authenticate(mock_client): return inner -def test_successful_copy(monkeypatch, powershell_exploiter): +def test_successful_copy(monkeypatch, powershell_exploiter, powershell_arguments): mock_client = MagicMock() mock_client.return_value.copy_file = MagicMock(return_value=True) monkeypatch.setattr(powershell, "PowerShellClient", mock_client) - success = powershell_exploiter.exploit_host() + exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) assert DROPPER_TARGET_PATH_64 in mock_client.return_value.copy_file.call_args[0][1] - assert success + assert exploit_result.exploitation_success -def test_failed_copy(monkeypatch, powershell_exploiter): +def test_failed_copy(monkeypatch, powershell_exploiter, powershell_arguments): mock_client = MagicMock() mock_client.return_value.copy_file = MagicMock(return_value=False) monkeypatch.setattr(powershell, "PowerShellClient", mock_client) - success = powershell_exploiter.exploit_host() - assert not success + exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) + assert not exploit_result.exploitation_success -def test_failed_monkey_execution(monkeypatch, powershell_exploiter): +def test_failed_monkey_execution(monkeypatch, powershell_exploiter, powershell_arguments): mock_client = MagicMock() mock_client.copy_file = MagicMock(return_value=True) mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception) @@ -144,11 +143,11 @@ def test_failed_monkey_execution(monkeypatch, powershell_exploiter): mock_powershell_client = MagicMock(side_effect=authenticate(mock_client)) monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client) - success = powershell_exploiter.exploit_host() - assert not success + exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) + assert not exploit_result.exploitation_success -def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter): +def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments): mock_client = MagicMock() mock_client.return_value.copy_file = MagicMock(return_value=True) @@ -159,7 +158,7 @@ def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter): monkeypatch.setattr(powershell, "PowerShellClient", mock_client) - powershell_exploiter.exploit_host() + powershell_exploiter.exploit_host(**powershell_arguments) # Total 6 attempts reported, 5 failed and 1 succeeded assert len(powershell_exploiter.exploit_attempts) == len(execute_cmd_returns)