From 1cb88e029addb8db30fa1a01f8b9fba83f9594de Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Tue, 4 Oct 2022 21:35:42 +0000
Subject: [PATCH] Agent: Extract method _exploit

---
 monkey/infection_monkey/exploit/sshexec.py | 62 ++++++++++++----------
 1 file changed, 35 insertions(+), 27 deletions(-)

diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py
index f6f703763..8fa5f9ad4 100644
--- a/monkey/infection_monkey/exploit/sshexec.py
+++ b/monkey/infection_monkey/exploit/sshexec.py
@@ -177,34 +177,12 @@ class SSHExploiter(HostExploiter):
         raise FailedExploitationError
 
     def _exploit_host(self) -> ExploiterResultData:
-        port = SSH_PORT
-
-        # if ssh banner found on different port, use that port.
-        for servkey, servdata in list(self.host.services.items()):
-            if servdata.get("name") == "ssh" and servkey.startswith("tcp-"):
-                port = int(servkey.replace("tcp-", ""))
-
-        is_open, _ = check_tcp_port(self.host.ip_addr, port)
-        if not is_open:
-            self.exploit_result.error_message = f"SSH port is closed on {self.host}, skipping"
-            self._publish_exploitation_event(
-                target=self.host.ip_addr,
-                exploitation_success=False,
-                error_message=self.exploit_result.error_message,
-                tags=(SSH_EXPLOITER_TAG,),
-            )
-            logger.info(self.exploit_result.error_message)
-            return self.exploit_result
-
         try:
-            ssh = self.exploit_with_ssh_keys(port)
-        except FailedExploitationError:
-            try:
-                ssh = self.exploit_with_login_creds(port)
-            except FailedExploitationError:
-                self.exploit_result.error_message = "Exploiter SSHExploiter is giving up..."
-                logger.error(self.exploit_result.error_message)
-                return self.exploit_result
+            ssh = self._exploit()
+        except FailedExploitationError as err:
+            self.exploit_result.error_message = str(err)
+            logger.error(str(err))
+            return self.exploit_result
 
         if self._is_interrupted():
             self._set_interrupted()
@@ -298,6 +276,36 @@ class SSHExploiter(HostExploiter):
             logger.error(self.exploit_result.error_message)
             return self.exploit_result
 
+    def _exploit(self) -> paramiko.SSHClient:
+        port = SSH_PORT
+
+        # if ssh banner found on different port, use that port.
+        for servkey, servdata in list(self.host.services.items()):
+            if servdata.get("name") == "ssh" and servkey.startswith("tcp-"):
+                port = int(servkey.replace("tcp-", ""))
+
+        is_open, _ = check_tcp_port(self.host.ip_addr, port)
+        if not is_open:
+            self.exploit_result.error_message = f"SSH port is closed on {self.host}, skipping"
+            self._publish_exploitation_event(
+                target=self.host.ip_addr,
+                exploitation_success=False,
+                error_message=self.exploit_result.error_message,
+                tags=(SSH_EXPLOITER_TAG,),
+            )
+            logger.info(self.exploit_result.error_message)
+            raise FailedExploitationError(self.exploit_result.error_message)
+
+        try:
+            ssh = self.exploit_with_ssh_keys(port)
+        except FailedExploitationError:
+            try:
+                ssh = self.exploit_with_login_creds(port)
+            except FailedExploitationError:
+                raise FailedExploitationError("Exploiter SSHExploiter is giving up...")
+
+        return ssh
+
     def _get_victim_os(self, ssh: paramiko.SSHClient) -> bool:
         try:
             _, stdout, _ = ssh.exec_command("uname -o", timeout=SSH_EXEC_TIMEOUT)