From 06899be264e89d0877eddafb956c0736653ce371 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 23 Mar 2022 13:25:22 -0400 Subject: [PATCH] Tests: Fix tests for PowerShellExploiter --- monkey/infection_monkey/exploit/powershell.py | 2 +- .../infection_monkey/exploit/conftest.py | 11 +-- .../exploit/test_powershell.py | 72 +++++-------------- 3 files changed, 26 insertions(+), 59 deletions(-) diff --git a/monkey/infection_monkey/exploit/powershell.py b/monkey/infection_monkey/exploit/powershell.py index 052b1f88f..3d5a41131 100644 --- a/monkey/infection_monkey/exploit/powershell.py +++ b/monkey/infection_monkey/exploit/powershell.py @@ -40,7 +40,7 @@ class PowerShellExploiter(HostExploiter): def _exploit_host(self): if not self._is_any_default_port_open(): - message = "No default PowerShell remoting ports are open." + message = "PowerShell Remoting appears to be disabled on the remote host" self.exploit_result.error_message = message logger.debug(message) diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/conftest.py b/monkey/tests/unit_tests/infection_monkey/exploit/conftest.py index 142a3065a..7d4265395 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/conftest.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/conftest.py @@ -21,8 +21,9 @@ def patch_win32api_get_user_name(local_user): sys.modules["win32api"] = win32api -def _create_host(http_enabled, https_enabled): +def _create_windows_host(http_enabled, https_enabled): host = MagicMock() + host.os = {"type": "windows"} host.services = {} if http_enabled: @@ -36,19 +37,19 @@ def _create_host(http_enabled, https_enabled): @pytest.fixture def https_only_host(): - return _create_host(False, True) + return _create_windows_host(False, True) @pytest.fixture def http_only_host(): - return _create_host(True, False) + return _create_windows_host(True, False) @pytest.fixture def http_and_https_both_enabled_host(): - return _create_host(True, True) + return _create_windows_host(True, True) @pytest.fixture def powershell_disabled_host(): - return _create_host(False, False) + return _create_windows_host(False, False) 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 c88ce99d7..698c7ac2d 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py @@ -5,8 +5,6 @@ from unittest.mock import MagicMock import pytest from infection_monkey.exploit import powershell -from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions -from infection_monkey.exploit.powershell_utils.credentials import Credentials from infection_monkey.model.host import VictimHost # Use the path_win32api_get_user_name fixture for all tests in this module @@ -19,19 +17,12 @@ NT_HASH_LIST = ["bogo_nt_1", "bogo_nt_2"] DROPPER_TARGET_PATH_64 = "C:\\agent64" -class AuthenticationErrorForTests(Exception): - pass - - mock_agent_repository = MagicMock() mock_agent_repository.get_agent_binary.return_value = BytesIO(b"BINARY_EXECUTABLE") -victim_host = VictimHost("127.0.0.1") -victim_host.os["type"] = "windows" - @pytest.fixture -def powershell_arguments(): +def powershell_arguments(http_and_https_both_enabled_host): options = { "dropper_target_path_win_64": DROPPER_TARGET_PATH_64, "credentials": { @@ -42,7 +33,7 @@ def powershell_arguments(): }, } arguments = { - "host": victim_host, + "host": http_and_https_both_enabled_host, "options": options, "current_depth": 2, "telemetry_messenger": MagicMock(), @@ -56,18 +47,13 @@ def powershell_arguments(): def powershell_exploiter(monkeypatch): pe = powershell.PowerShellExploiter() - monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests) monkeypatch.setattr(powershell, "is_windows_os", lambda: True) return pe -def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_arguments): - mock_powershell_client = MagicMock() - mock_powershell_client.connect = MagicMock(side_effect=Exception) - monkeypatch.setattr( - powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) - ) +def test_powershell_disabled(powershell_exploiter, powershell_arguments, powershell_disabled_host): + powershell_arguments["host"] = powershell_disabled_host exploit_result = powershell_exploiter.exploit_host(**powershell_arguments) assert not exploit_result.exploitation_success @@ -75,15 +61,10 @@ def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_argum assert "disabled" in exploit_result.error_message -def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments): - def allow_http(_, credentials: Credentials, auth_options: AuthOptions): - if not auth_options.ssl: - raise AuthenticationErrorForTests - else: - raise Exception +def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments, http_only_host): + powershell_arguments["host"] = http_only_host mock_powershell_client = MagicMock() - mock_powershell_client.connect = MagicMock(side_effect=allow_http) monkeypatch.setattr( powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) ) @@ -94,29 +75,26 @@ def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments assert not call_args[0][2].ssl -def test_powershell_https(monkeypatch, powershell_exploiter, powershell_arguments): - def allow_https(_, credentials: Credentials, auth_options: AuthOptions): - if auth_options.ssl: - raise AuthenticationErrorForTests - else: - raise Exception +def test_powershell_https(monkeypatch, powershell_exploiter, powershell_arguments, https_only_host): + powershell_arguments["host"] = https_only_host mock_powershell_client = MagicMock() - mock_powershell_client.connect = MagicMock(side_effect=allow_https) - monkeypatch.setattr( - powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) - ) + mock_powershell_client.connect = MagicMock(side_effect=Exception("Failed login")) + mock_powershell_client_constructor = MagicMock(return_value=mock_powershell_client) + monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client_constructor) 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": + for call_args in mock_powershell_client_constructor.call_args_list: + if call_args[0][1].secret != "": assert call_args[0][2].ssl + else: + assert not call_args[0][2].ssl def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_arguments): mock_powershell_client = MagicMock() - mock_powershell_client.connect = MagicMock(side_effect=AuthenticationErrorForTests) + mock_powershell_client.connect = MagicMock(side_effect=Exception("Failed login")) monkeypatch.setattr( powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client) ) @@ -127,16 +105,6 @@ def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_argu assert "Unable to authenticate" in exploit_result.error_message -def authenticate(mock_client): - def inner(_, credentials: Credentials, auth_options: AuthOptions): - if credentials.username == "user1" and credentials.secret == "pass2": - return mock_client - else: - raise AuthenticationErrorForTests("Invalid credentials") - - return inner - - def test_successful_copy(monkeypatch, powershell_exploiter, powershell_arguments): mock_client = MagicMock() @@ -188,11 +156,9 @@ def test_successful_propagation(monkeypatch, powershell_exploiter, powershell_ar def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments): - # 1st call is for determining HTTP/HTTPs. 6 remaining calls are actual login attempts. the 6th - # login attempt doesn't throw an exception, signifying that login with credentials was - # successful. - connection_attempts = [True, Exception, Exception, Exception, Exception, Exception, True] - mock_powershell_client = MagicMock(side_effect=connection_attempts) + # First 5 login attempts fail. The 6th is successful. + connection_attempts = [Exception, Exception, Exception, Exception, Exception, True] + mock_powershell_client = MagicMock() mock_powershell_client.connect = MagicMock(side_effect=connection_attempts) monkeypatch.setattr( powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)