Agent: Stop sending PropagationEvent before attempt

monkey/infection_monkey/exploit/sshexec.py

@@ -2,6 +2,7 @@ import io
 import logging
 from ipaddress import IPv4Address
 from pathlib import PurePath
+from typing import Optional
 
 import paramiko
 
@@ -19,6 +20,7 @@ from common.types import PortStatus
 from common.utils import Timer
 from common.utils.attack_utils import ScanStatus
 from common.utils.exceptions import FailedExploitationError
+from infection_monkey.exploit import RetrievalError
 from infection_monkey.exploit.HostExploiter import HostExploiter
 from infection_monkey.exploit.tools.helpers import get_agent_dst_path
 from infection_monkey.i_puppet import ExploiterResultData
@@ -205,14 +207,19 @@ class SSHExploiter(HostExploiter):
             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
+        except RuntimeError as err:
+            error_message = str(err)
+            self.exploit_result.error_message = error_message
+            logger.error(error_message)
+        finally:
+            ssh.close()
+            return self.exploit_result
 
     def _exploit(self, port) -> paramiko.SSHClient:
         try:
@@ -226,23 +233,13 @@ class SSHExploiter(HostExploiter):
         return ssh
 
     def _propagate(self, ssh: paramiko.SSHClient):
-        if not self.host.os.get("type") and not self._get_victim_os(ssh):
+        agent_binary_file_object = self._get_agent_binary(ssh)
-            raise FailedExploitationError(
+        if agent_binary_file_object is None:
-                f"Can't find suitable monkey executable for host {self.host}"
+            raise RuntimeError("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:
-            raise FailedExploitationError(
-                f"Can't find suitable monkey executable for host {self.host}"
-            )
 
         if self._is_interrupted():
             self._set_interrupted()
-            raise FailedExploitationError(f"Propagation was interrupted")
+            raise RuntimeError("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)
@@ -327,6 +324,19 @@ class SSHExploiter(HostExploiter):
             return False
         return True
 
+    def _get_agent_binary(self, ssh: paramiko.SSHClient) -> Optional[io.BytesIO]:
+        if not self.host.os.get("type") and not self._get_victim_os(ssh):
+            return None
+
+        try:
+            agent_binary_file_object = self.agent_binary_repository.get_agent_binary(
+                self.exploit_result.os
+            )
+        except RetrievalError:
+            return None
+
+        return agent_binary_file_object
+
     def _upload_agent_binary(
         self,
         ssh: paramiko.SSHClient,