Agent: Add connect() method to IPowerShellClient

monkey/infection_monkey/exploit/powershell.py

@@ -127,9 +127,10 @@ class PowerShellExploiter(HostExploiter):
             encryption=ENCRYPTION_AUTO,
             ssl=use_ssl,
         )
-        # TODO: Report login attempt or find a better way of detecting if SSL is enabled
 
-        PowerShellClient(self.host.ip_addr, credentials, auth_options)
+        # TODO: Report login attempt or find a better way of detecting if SSL is enabled
+        client = PowerShellClient(self.host.ip_addr, credentials, auth_options)
+        client.connect()
 
     def _authenticate_via_brute_force(
         self, credentials: List[Credentials], auth_options: List[AuthOptions]
@@ -137,6 +138,7 @@ class PowerShellExploiter(HostExploiter):
         for (creds, opts) in zip(credentials, auth_options):
             try:
                 client = PowerShellClient(self.host.ip_addr, creds, opts)
+                client.connect()
                 logger.info(
                     f"Successfully logged into {self.host.ip_addr} using Powershell. User: "
                     f"{creds.username}, Secret Type: {creds.secret_type.name}"

monkey/infection_monkey/exploit/powershell_utils/powershell_client.py

@@ -54,6 +54,10 @@ def format_password(credentials: Credentials) -> Optional[str]:
 
 
 class IPowerShellClient(Protocol, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def connect(self) -> str:
+        pass
+
     @abc.abstractmethod
     def execute_cmd(self, cmd: str) -> str:
         pass
@@ -72,14 +76,19 @@ class PowerShellClient(IPowerShellClient):
         _set_sensitive_packages_log_level_to_error()
 
         self._ip_addr = ip_addr
+        self._credentials = credentials
+        self._auth_options = auth_options
+        self._client = None
+
+    def connect(self):
         self._client = Client(
-            ip_addr,
+            self._ip_addr,
-            username=credentials.username,
+            username=self._credentials.username,
-            password=format_password(credentials),
+            password=format_password(self._credentials),
             cert_validation=False,
-            auth=auth_options.auth_type,
+            auth=self._auth_options.auth_type,
-            encryption=auth_options.encryption,
+            encryption=self._auth_options.encryption,
-            ssl=auth_options.ssl,
+            ssl=self._auth_options.ssl,
             connection_timeout=CONNECTION_TIMEOUT,
         )
 

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

@@ -58,8 +58,11 @@ def powershell_exploiter(monkeypatch):
 
 
 def test_powershell_disabled(monkeypatch, powershell_exploiter, powershell_arguments):
-    mock_powershell_client = MagicMock(side_effect=Exception)
+    mock_powershell_client = MagicMock()
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    mock_powershell_client.connect = MagicMock(side_effect=Exception)
+    monkeypatch.setattr(
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
 
     exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
     assert not exploit_result.exploitation_success
@@ -74,8 +77,12 @@ def test_powershell_http(monkeypatch, powershell_exploiter, powershell_arguments
         else:
             raise Exception
 
-    mock_powershell_client = MagicMock(side_effect=allow_http)
+    mock_powershell_client = MagicMock()
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    mock_powershell_client.connect = MagicMock(side_effect=allow_http)
+    monkeypatch.setattr(
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
+
     powershell_exploiter.exploit_host(**powershell_arguments)
 
     for call_args in mock_powershell_client.call_args_list:
@@ -89,8 +96,12 @@ def test_powershell_https(monkeypatch, powershell_exploiter, powershell_argument
         else:
             raise Exception
 
-    mock_powershell_client = MagicMock(side_effect=allow_https)
+    mock_powershell_client = MagicMock()
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    mock_powershell_client.connect = MagicMock(side_effect=allow_https)
+    monkeypatch.setattr(
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
+
     powershell_exploiter.exploit_host(**powershell_arguments)
 
     for call_args in mock_powershell_client.call_args_list:
@@ -99,8 +110,11 @@ def test_powershell_https(monkeypatch, powershell_exploiter, powershell_argument
 
 
 def test_no_valid_credentials(monkeypatch, powershell_exploiter, powershell_arguments):
-    mock_powershell_client = MagicMock(side_effect=AuthenticationErrorForTests)
+    mock_powershell_client = MagicMock()
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    mock_powershell_client.connect = MagicMock(side_effect=AuthenticationErrorForTests)
+    monkeypatch.setattr(
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
 
     exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
     assert not exploit_result.exploitation_success
@@ -142,13 +156,14 @@ def test_failed_copy(monkeypatch, powershell_exploiter, powershell_arguments):
 
 
 def test_failed_monkey_execution(monkeypatch, powershell_exploiter, powershell_arguments):
-    mock_client = MagicMock()
+    mock_powershell_client = MagicMock()
-    mock_client.execute_cmd_as_detached_process = MagicMock(
+    mock_powershell_client.execute_cmd_as_detached_process = MagicMock(
         side_effect=Exception("EXECUTION FAILED")
     )
 
-    mock_powershell_client = MagicMock(side_effect=authenticate(mock_client))
+    monkeypatch.setattr(
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
 
     exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)
     assert exploit_result.exploitation_success is True
@@ -172,8 +187,11 @@ def test_login_attempts_correctly_reported(monkeypatch, powershell_exploiter, po
     # login attempt doesn't throw an exception, signifying that login with credentials was
     # successful.
     connection_attempts = [True, Exception, Exception, Exception, Exception, Exception, True]
-    mock_client = MagicMock(side_effect=connection_attempts)
+    mock_powershell_client = MagicMock(side_effect=connection_attempts)
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
+    mock_powershell_client.connect = MagicMock(side_effect=connection_attempts)
+    monkeypatch.setattr(
+        powershell, "PowerShellClient", MagicMock(return_value=mock_powershell_client)
+    )
 
     exploit_result = powershell_exploiter.exploit_host(**powershell_arguments)