Tests: Add HTTP vs HTPS unit tests for PowerShellExploiter

monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py

@@ -0,0 +1,68 @@
+from collections import namedtuple
+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
+
+USER_LIST = ["user1", "user2"]
+PASSWORD_LIST = ["pass1", "pass2"]
+
+Config = namedtuple("Config", ["exploit_user_list", "exploit_password_list"])
+
+
+class TestAuthenticationError(Exception):
+    pass
+
+
+@pytest.fixture
+def powershell_exploiter(monkeypatch):
+    host = VictimHost("127.0.0.1")
+    pe = powershell.PowerShellExploiter(host)
+    pe._config = Config(USER_LIST, PASSWORD_LIST)
+
+    monkeypatch.setattr(powershell, "AuthenticationError", TestAuthenticationError)
+
+    return pe
+
+
+def test_powershell_disabled(monkeypatch, powershell_exploiter):
+    mock_powershell_client = MagicMock(side_effect=Exception)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+
+    success = powershell_exploiter.exploit_host()
+    assert not success
+
+
+def test_powershell_http(monkeypatch, powershell_exploiter):
+    def allow_http(_, credentials: Credentials, auth_options: AuthOptions):
+        if not auth_options.ssl:
+            raise TestAuthenticationError
+        else:
+            raise Exception
+
+    mock_powershell_client = MagicMock(side_effect=allow_http)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    powershell_exploiter.exploit_host()
+
+    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 allow_https(_, credentials: Credentials, auth_options: AuthOptions):
+        if auth_options.ssl:
+            raise TestAuthenticationError
+        else:
+            raise Exception
+
+    mock_powershell_client = MagicMock(side_effect=allow_https)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    powershell_exploiter.exploit_host()
+
+    for call_args in mock_powershell_client.call_args_list:
+        if call_args[0][1].password != "" and call_args[0][1].password != "dummy_password":
+            assert call_args[0][2].ssl