UT: Add arguments and return exploit result data to PowerShell exploit

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

@@ -1,4 +1,3 @@
-from collections import namedtuple
 from unittest.mock import MagicMock
 
 import pytest
@@ -15,57 +14,57 @@ USER_LIST = ["user1", "user2"]
 PASSWORD_LIST = ["pass1", "pass2"]
 LM_HASH_LIST = ["bogo_lm_1"]
 NT_HASH_LIST = ["bogo_nt_1", "bogo_nt_2"]
-DROPPER_TARGET_PATH_32 = "C:\\agent32"
 DROPPER_TARGET_PATH_64 = "C:\\agent64"
 
-Config = namedtuple(
-    "Config",
-    [
-        "exploit_user_list",
-        "exploit_password_list",
-        "exploit_lm_hash_list",
-        "exploit_ntlm_hash_list",
-        "dropper_target_path_win_64",
-    ],
-)
-
 
 class AuthenticationErrorForTests(Exception):
     pass
 
 
+@pytest.fixture
+def powershell_arguments():
+    options = {
+        "dropper_target_path_win_64": DROPPER_TARGET_PATH_64,
+        "credentials": {
+            "exploit_user_list": USER_LIST,
+            "exploit_password_list": PASSWORD_LIST,
+            "exploit_lm_hash_list": LM_HASH_LIST,
+            "exploit_ntlm_hash_list": NT_HASH_LIST,
+        },
+    }
+    arguments = {
+        "host": VictimHost("127.0.0.1"),
+        "options": options,
+        "current_depth": 2,
+        "telemetry_messenger": MagicMock(),
+        "agent_repository": MagicMock(),
+    }
+    return arguments
+
+
 @pytest.fixture
 def powershell_exploiter(monkeypatch):
-    host = VictimHost("127.0.0.1")
+    pe = powershell.PowerShellExploiter()
-    pe = powershell.PowerShellExploiter(host)
-    pe._config = Config(
-        USER_LIST,
-        PASSWORD_LIST,
-        LM_HASH_LIST,
-        NT_HASH_LIST,
-        DROPPER_TARGET_PATH_32,
-        DROPPER_TARGET_PATH_64,
-    )
 
     monkeypatch.setattr(powershell, "AuthenticationError", AuthenticationErrorForTests)
     monkeypatch.setattr(powershell, "is_windows_os", lambda: True)
     # It's regrettable to mock out a private method on the PowerShellExploiter instance object, but
     # it's necessary to avoid having to deal with the monkeyfs. TODO: monkeyfs has been removed, so
     # fix this.
-    monkeypatch.setattr(pe, "_write_virtual_file_to_local_path", lambda: None)
+    # monkeypatch.setattr(pe, "_create_local_agent_file", lambda: None)
 
     return pe
 
 
-def test_powershell_disabled(monkeypatch, powershell_exploiter):
+def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_powershell_client = MagicMock(side_effect=Exception)
     monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
 
-    success = powershell_exploiter.exploit_host()
+    exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
-    assert not success
+    assert not exploit_result.exploitation_success
 
 
-def test_powershell_http(monkeypatch, powershell_exploiter):
+def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments):
     def allow_http(_, credentials: Credentials, auth_options: AuthOptions):
         if not auth_options.ssl:
             raise AuthenticationErrorForTests
@@ -74,13 +73,13 @@ def test_powershell_http(monkeypatch, powershell_exploiter):
 
     mock_powershell_client = MagicMock(side_effect=allow_http)
     monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
-    powershell_exploiter.exploit_host()
+    powershell_exploiter.exploit_host(**powershell_arguments)
 
     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 test_powershell_https(monkeypatch, powershell_exploiter, powershell_arguments):
     def allow_https(_, credentials: Credentials, auth_options: AuthOptions):
         if auth_options.ssl:
             raise AuthenticationErrorForTests
@@ -89,19 +88,19 @@ def test_powershell_https(monkeypatch, powershell_exploiter):
 
     mock_powershell_client = MagicMock(side_effect=allow_https)
     monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
-    powershell_exploiter.exploit_host()
+    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":
             assert call_args[0][2].ssl
 
 
-def test_no_valid_credentials(monkeypatch, powershell_exploiter):
+def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_powershell_client = MagicMock(side_effect=AuthenticationErrorForTests)
     monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
 
-    success = powershell_exploiter.exploit_host()
+    exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
-    assert not success
+    assert not exploit_result.exploitation_success
 
 
 def authenticate(mock_client):
@@ -114,29 +113,29 @@ def authenticate(mock_client):
     return inner
 
 
-def test_successful_copy(monkeypatch, powershell_exploiter):
+def test_successful_copy(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_client = MagicMock()
     mock_client.return_value.copy_file = MagicMock(return_value=True)
 
     monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
 
-    success = powershell_exploiter.exploit_host()
+    exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
 
     assert DROPPER_TARGET_PATH_64 in mock_client.return_value.copy_file.call_args[0][1]
-    assert success
+    assert exploit_result.exploitation_success
 
 
-def test_failed_copy(monkeypatch, powershell_exploiter):
+def test_failed_copy(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_client = MagicMock()
     mock_client.return_value.copy_file = MagicMock(return_value=False)
 
     monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
 
-    success = powershell_exploiter.exploit_host()
+    exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
-    assert not success
+    assert not exploit_result.exploitation_success
 
 
-def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
+def test_failed_monkey_execution(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_client = MagicMock()
     mock_client.copy_file = MagicMock(return_value=True)
     mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception)
@@ -144,11 +143,11 @@ def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
     mock_powershell_client = MagicMock(side_effect=authenticate(mock_client))
     monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
 
-    success = powershell_exploiter.exploit_host()
+    exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
-    assert not success
+    assert not exploit_result.exploitation_success
 
 
-def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter):
+def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter, powershell_arguments):
     mock_client = MagicMock()
     mock_client.return_value.copy_file = MagicMock(return_value=True)
 
@@ -159,7 +158,7 @@ def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter):
 
     monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
 
-    powershell_exploiter.exploit_host()
+    powershell_exploiter.exploit_host(**powershell_arguments)
 
     # Total 6 attempts reported, 5 failed and 1 succeeded
     assert len(powershell_exploiter.exploit_attempts) == len(execute_cmd_returns)