Tests: Fix tests for PowerShellExploiter

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)
 

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)

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)