Tests: Fix unit tests for powershell_utils.auth_options

monkey/tests/unit_tests/infection_monkey/exploit/powershell_utils/test_auth_options.py

@@ -1,3 +1,7 @@
+from unittest.mock import MagicMock
+
+import pytest
+
 # from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions
 from infection_monkey.exploit.powershell_utils.auth_options import (
     AUTH_BASIC,
@@ -16,85 +20,124 @@ CREDENTIALS_LM_HASH = Credentials("user4", "LM_HASH:NONE", SecretType.LM_HASH)
 CREDENTIALS_NT_HASH = Credentials("user5", "NONE:NT_HASH", SecretType.NT_HASH)
 
 
-def test_get_auth_options__ssl_true_with_password():
-    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, use_ssl=True)
+def _create_host(http_enabled, https_enabled):
+    host = MagicMock()
+    host.services = {}
+
+    if http_enabled:
+        host.services["tcp-5985"] = {}
+
+    if https_enabled:
+        host.services["tcp-5986"] = {}
+
+    return host
+
+
+@pytest.fixture
+def https_only_host():
+    return _create_host(False, True)
+
+
+@pytest.fixture
+def http_only_host():
+    return _create_host(True, False)
+
+
+@pytest.fixture
+def http_and_https_both_enabled_host():
+    return _create_host(True, True)
+
+
+@pytest.fixture
+def powershell_disabled_host():
+    return _create_host(False, False)
+
+
+def test_get_auth_options__ssl_true_with_password(https_only_host):
+    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, https_only_host)
 
     assert auth_options.ssl
 
 
-def test_get_auth_options__ssl_true_empty_password():
-    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, use_ssl=True)
-
-    assert not auth_options.ssl
-
-
-def test_get_auth_options__ssl_true_none_password():
-    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, use_ssl=True)
+def test_get_auth_options__ssl_preferred(http_and_https_both_enabled_host):
+    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, http_and_https_both_enabled_host)
 
     assert auth_options.ssl
 
 
-def test_get_auth_options__ssl_false_with_password():
-    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, use_ssl=False)
+def test_get_auth_options__ssl_true_empty_password(https_only_host):
+    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, https_only_host)
 
     assert not auth_options.ssl
 
 
-def test_get_auth_options__ssl_false_empty_password():
-    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, use_ssl=False)
+def test_get_auth_options__ssl_true_none_password(https_only_host):
+    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, https_only_host)
+
+    assert auth_options.ssl
+
+
+def test_get_auth_options__ssl_false_with_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, http_only_host)
 
     assert not auth_options.ssl
 
 
-def test_get_auth_options__ssl_false_none_password():
-    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, use_ssl=False)
+def test_get_auth_options__ssl_false_empty_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, http_only_host)
 
     assert not auth_options.ssl
 
 
-def test_get_auth_options__auth_type_with_password():
-    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, use_ssl=False)
+def test_get_auth_options__ssl_false_none_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, http_only_host)
+
+    assert not auth_options.ssl
+
+
+def test_get_auth_options__auth_type_with_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, http_only_host)
 
     assert auth_options.auth_type == AUTH_NEGOTIATE
 
 
-def test_get_auth_options__auth_type_empty_password():
-    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, use_ssl=False)
+def test_get_auth_options__auth_type_empty_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, http_only_host)
 
     assert auth_options.auth_type == AUTH_BASIC
 
 
-def test_get_auth_options__auth_type_none_password():
-    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, use_ssl=False)
+def test_get_auth_options__auth_type_none_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, http_only_host)
 
     assert auth_options.auth_type == AUTH_NEGOTIATE
 
 
-def test_get_auth_options__auth_type_with_LM_hash():
-    auth_options = get_auth_options(CREDENTIALS_LM_HASH, use_ssl=False)
+def test_get_auth_options__auth_type_with_LM_hash(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_LM_HASH, http_only_host)
 
     assert auth_options.auth_type == AUTH_NTLM
 
 
-def test_get_auth_options__auth_type_with_NT_hash():
-    auth_options = get_auth_options(CREDENTIALS_NT_HASH, use_ssl=False)
+def test_get_auth_options__auth_type_with_NT_hash(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_NT_HASH, http_only_host)
 
     assert auth_options.auth_type == AUTH_NTLM
 
 
-def test_get_auth_options__encryption_with_password():
-    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, use_ssl=False)
+def test_get_auth_options__encryption_with_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_WITH_PASSWORD, http_only_host)
 
     assert auth_options.encryption == ENCRYPTION_AUTO
 
 
-def test_get_auth_options__encryption_empty_password():
-    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, use_ssl=False)
+def test_get_auth_options__encryption_empty_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_EMPTY_PASSWORD, http_only_host)
 
     assert auth_options.encryption == ENCRYPTION_NEVER
 
 
-def test_get_auth_options__encryption_none_password():
-    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, use_ssl=False)
+def test_get_auth_options__encryption_none_password(http_only_host):
+    auth_options = get_auth_options(CREDENTIALS_NONE_PASSWORD, http_only_host)
 
     assert auth_options.encryption == ENCRYPTION_AUTO