From 0a1901b9a1ba24ad8033d815a2f59422e83fd969 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Tue, 4 Oct 2022 22:12:50 +0000
Subject: [PATCH] Agent: Use error to propagate failure

---
 monkey/infection_monkey/exploit/sshexec.py | 57 ++++++++--------------
 1 file changed, 21 insertions(+), 36 deletions(-)

diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py
index c8ea2f415..31757a5e3 100644
--- a/monkey/infection_monkey/exploit/sshexec.py
+++ b/monkey/infection_monkey/exploit/sshexec.py
@@ -188,7 +188,20 @@ class SSHExploiter(HostExploiter):
             self._set_interrupted()
             return self.exploit_result
 
-        return self._propagate(ssh)
+        try:
+            self._propagate(ssh)
+        except FailedExploitationError as err:
+            ssh.close()
+            self.exploit_result.error_message = str(err)
+            logger.error(self.exploit_result.error_message)
+
+            self._publish_propagation_event(
+                target=self.host.ip_addr,
+                propagation_success=False,
+                error_message=self.exploit_result.error_message,
+                tags=PROPAGATION_TAGS,
+            )
+        return self.exploit_result
 
     def _exploit(self) -> paramiko.SSHClient:
         port = SSH_PORT
@@ -222,33 +235,24 @@ class SSHExploiter(HostExploiter):
 
     def _propagate(self, ssh: paramiko.SSHClient):
         if not self.host.os.get("type") and not self._get_victim_os(ssh):
-            return self.exploit_result
+            raise FailedExploitationError(
+                f"Can't find suitable monkey executable for host {self.host}"
+            )
 
         agent_binary_file_object = self.agent_binary_repository.get_agent_binary(
             self.exploit_result.os
         )
 
         if not agent_binary_file_object:
-            self.exploit_result.error_message = (
+            raise FailedExploitationError(
                 f"Can't find suitable monkey executable for host {self.host}"
             )
 
-            self._publish_propagation_event(
-                target=self.host.ip_addr,
-                propagation_success=False,
-                error_message=self.exploit_result.error_message,
-                tags=(SSH_EXPLOITER_TAG,),
-            )
-
-            logger.error(self.exploit_result.error_message)
-            return self.exploit_result
-
         if self._is_interrupted():
             self._set_interrupted()
-            return self.exploit_result
+            raise FailedExploitationError(f"Propagation was interrupted")
 
         monkey_path_on_victim = get_agent_dst_path(self.host)
-
         status = self._upload_agent_binary(ssh, agent_binary_file_object, monkey_path_on_victim)
 
         self.telemetry_messenger.send_telemetry(
@@ -261,13 +265,7 @@ class SSHExploiter(HostExploiter):
         )
 
         if status == ScanStatus.SCANNED:
-            self._publish_propagation_event(
-                target=self.host.ip_addr,
-                propagation_success=False,
-                error_message=self.exploit_result.error_message,
-                tags=PROPAGATION_TAGS,
-            )
-            return self.exploit_result
+            raise FailedExploitationError(self.exploit_result.error_message)
 
         try:
             cmdline = f"{monkey_path_on_victim} {MONKEY_ARG}"
@@ -292,22 +290,9 @@ class SSHExploiter(HostExploiter):
 
             ssh.close()
             self.add_executed_cmd(cmdline)
-            return self.exploit_result
 
         except Exception as exc:
-            self.exploit_result.error_message = (
-                f"Error running monkey on victim {self.host}: ({exc})"
-            )
-
-            self._publish_propagation_event(
-                target=self.host.ip_addr,
-                propagation_success=False,
-                error_message=self.exploit_result.error_message,
-                tags=PROPAGATION_TAGS,
-            )
-
-            logger.error(self.exploit_result.error_message)
-            return self.exploit_result
+            raise FailedExploitationError(f"Error running monkey on victim {self.host}: ({exc})")
 
     def _get_victim_os(self, ssh: paramiko.SSHClient) -> bool:
         try: