Tests: Add more tests for PowerShellExploiter

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

@@ -4,14 +4,25 @@ from unittest.mock import MagicMock
 import pytest
 
 from infection_monkey.exploit import powershell
+from infection_monkey.exploit.consts import WIN_ARCH_32, WIN_ARCH_64
 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"]
+DROPPER_TARGET_PATH_32 = "C:\\agent32"
+DROPPER_TARGET_PATH_64 = "C:\\agent64"
 
-Config = namedtuple("Config", ["exploit_user_list", "exploit_password_list"])
+Config = namedtuple(
+    "Config",
+    [
+        "exploit_user_list",
+        "exploit_password_list",
+        "dropper_target_path_win_32",
+        "dropper_target_path_win_64",
+    ],
+)
 
 
 class TestAuthenticationError(Exception):
@@ -22,9 +33,12 @@ class TestAuthenticationError(Exception):
 def powershell_exploiter(monkeypatch):
     host = VictimHost("127.0.0.1")
     pe = powershell.PowerShellExploiter(host)
-    pe._config = Config(USER_LIST, PASSWORD_LIST)
+    pe._config = Config(USER_LIST, PASSWORD_LIST, DROPPER_TARGET_PATH_32, DROPPER_TARGET_PATH_64)
 
     monkeypatch.setattr(powershell, "AuthenticationError", TestAuthenticationError)
+    # 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
+    monkeypatch.setattr(pe, "_write_virtual_file_to_local_path", lambda: None)
 
     return pe
 
@@ -66,3 +80,64 @@ def test_powershell_https(monkeypatch, powershell_exploiter):
     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
+
+
+def test_no_valid_credentials(monkeypatch, powershell_exploiter):
+    mock_powershell_client = MagicMock(side_effect=TestAuthenticationError)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+
+    success = powershell_exploiter.exploit_host()
+    assert not success
+
+
+def authenticate(mock_client):
+    def inner(_, credentials: Credentials, auth_options: AuthOptions):
+        if credentials.username == "user1" and credentials.password == "pass2":
+            return mock_client
+        else:
+            raise TestAuthenticationError("Invalid credentials")
+
+    return inner
+
+
+@pytest.mark.parametrize(
+    "dropper_target_path,arch",
+    [(DROPPER_TARGET_PATH_32, WIN_ARCH_32), (DROPPER_TARGET_PATH_64, WIN_ARCH_64)],
+)
+def test_successful_copy(monkeypatch, powershell_exploiter, dropper_target_path, arch):
+    mock_client = MagicMock()
+    mock_client.get_host_architecture = lambda: arch
+    mock_client.copy_file = MagicMock(return_value=True)
+
+    mock_powershell_client = MagicMock(side_effect=authenticate(mock_client))
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+
+    success = powershell_exploiter.exploit_host()
+
+    assert dropper_target_path in mock_client.copy_file.call_args[0][1]
+    assert success
+
+
+def test_failed_copy(monkeypatch, powershell_exploiter):
+    mock_client = MagicMock()
+    mock_client.get_host_architecture = lambda: WIN_ARCH_32
+    mock_client.copy_file = MagicMock(return_value=False)
+
+    mock_powershell_client = MagicMock(side_effect=authenticate(mock_client))
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+
+    success = powershell_exploiter.exploit_host()
+    assert not success
+
+
+def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
+    mock_client = MagicMock()
+    mock_client.get_host_architecture = lambda: WIN_ARCH_32
+    mock_client.copy_file = MagicMock(return_value=True)
+    mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception)
+
+    mock_powershell_client = MagicMock(side_effect=authenticate(mock_client))
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+
+    success = powershell_exploiter.exploit_host()
+    assert not success