From 3bca02af59cc188b64758c828b791af93b89c742 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 5 Oct 2022 14:25:02 +0000
Subject: [PATCH] Agent: Fix powershell tests

---
 .../exploit/test_powershell.py                | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

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 97096830d..5fb98ea93 100644
--- a/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
+++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py
@@ -16,6 +16,7 @@ LM_HASH_LIST = ["bogo_lm_1"]
 NT_HASH_LIST = ["bogo_nt_1", "bogo_nt_2"]
 
 bogus_servers = ["1.1.1.1:5000", "2.2.2.2:5007"]
+VICTIM_IP = "10.10.10.1"
 
 
 mock_agent_binary_repository = MagicMock()
@@ -23,7 +24,25 @@ mock_agent_binary_repository.get_agent_binary.return_value = BytesIO(b"BINARY_EX
 
 
 @pytest.fixture
-def powershell_arguments(http_and_https_both_enabled_host):
+def host_with_ip_address(http_and_https_both_enabled_host):
+    http_and_https_both_enabled_host.ip_addr = VICTIM_IP
+    return http_and_https_both_enabled_host
+
+
+@pytest.fixture
+def http_host_with_ip_address(http_only_host):
+    http_only_host.ip_addr = VICTIM_IP
+    return http_only_host
+
+
+@pytest.fixture
+def https_host_with_ip_address(https_only_host):
+    https_only_host.ip_addr = VICTIM_IP
+    return https_only_host
+
+
+@pytest.fixture
+def powershell_arguments(host_with_ip_address):
     options = {
         "credentials": {
             "exploit_user_list": USER_LIST,
@@ -33,7 +52,7 @@ def powershell_arguments(http_and_https_both_enabled_host):
         },
     }
     arguments = {
-        "host": http_and_https_both_enabled_host,
+        "host": host_with_ip_address,
         "servers": bogus_servers,
         "options": options,
         "current_depth": 2,
@@ -63,8 +82,10 @@ def test_powershell_disabled(powershell_exploiter, powershell_arguments, powersh
     assert "disabled" in exploit_result.error_message
 
 
-def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments, http_only_host):
-    powershell_arguments["host"] = http_only_host
+def test_powershell_http(
+    monkeypatch, powershell_exploiter, powershell_arguments, http_host_with_ip_address
+):
+    powershell_arguments["host"] = http_host_with_ip_address
 
     mock_powershell_client = MagicMock()
     monkeypatch.setattr(
@@ -77,7 +98,7 @@ 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, https_only_host):
+def test_powershell_https(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_powershell_client = MagicMock()
     mock_powershell_client.connect = MagicMock(side_effect=Exception("Failed login"))
     mock_powershell_client_constructor = MagicMock(return_value=mock_powershell_client)
@@ -191,11 +212,11 @@ def test_build_monkey_execution_command():
 
 
 def test_skip_http_only_logins(
-    monkeypatch, powershell_exploiter, powershell_arguments, https_only_host
+    monkeypatch, powershell_exploiter, powershell_arguments, https_host_with_ip_address
 ):
     # Only HTTPS is enabled on the destination, so we should never try to connect with "" empty
     # password, since connection with empty password requires SSL == False.
-    powershell_arguments["host"] = https_only_host
+    powershell_arguments["host"] = https_host_with_ip_address
 
     mock_powershell_client = MagicMock()
     mock_powershell_client.connect = MagicMock(side_effect=Exception("Failed login"))