From e11bd2c7f2aa5c9d31b9b1d2b006b671cdcbb846 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 5 Oct 2022 20:09:32 +0000
Subject: [PATCH] Agent: Stamp start time prior to running exploit

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

diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py
index c0b934bcd..563b85ac6 100644
--- a/monkey/infection_monkey/exploit/sshexec.py
+++ b/monkey/infection_monkey/exploit/sshexec.py
@@ -2,6 +2,7 @@ import io
 import logging
 from ipaddress import IPv4Address
 from pathlib import PurePath
+from time import time
 from typing import Optional, Tuple
 
 import paramiko
@@ -89,6 +90,8 @@ class SSHExploiter(HostExploiter):
                 pkey = paramiko.RSAKey.from_private_key(pkey)
             except (IOError, paramiko.SSHException, paramiko.PasswordRequiredException):
                 logger.error("Failed reading ssh key")
+
+            stamp = time()
             try:
                 ssh.connect(
                     self.host.ip_addr,
@@ -105,7 +108,7 @@ class SSHExploiter(HostExploiter):
                 )
                 self.add_vuln_port(port)
                 self.exploit_result.exploitation_success = True
-                self._publish_exploitation_event(True)
+                self._publish_exploitation_event(stamp, True)
                 self.report_login_attempt(True, user, ssh_key=ssh_string)
                 return ssh
             except paramiko.AuthenticationException as err:
@@ -114,7 +117,7 @@ class SSHExploiter(HostExploiter):
                     f"Failed logging into victim {self.host} with {ssh_string} private key: {err}"
                 )
                 logger.info(error_message)
-                self._publish_exploitation_event(False, error_message=error_message)
+                self._publish_exploitation_event(stamp, False, error_message=error_message)
                 self.report_login_attempt(False, user, ssh_key=ssh_string)
                 continue
             except Exception as err:
@@ -138,6 +141,8 @@ class SSHExploiter(HostExploiter):
 
             ssh = paramiko.SSHClient()
             ssh.set_missing_host_key_policy(paramiko.WarningPolicy())
+
+            stamp = time()
             try:
                 ssh.connect(
                     self.host.ip_addr,
@@ -153,14 +158,14 @@ 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._publish_exploitation_event(True)
+                self._publish_exploitation_event(stamp, True)
                 self.report_login_attempt(True, user, current_password)
                 return ssh
 
             except paramiko.AuthenticationException as err:
                 error_message = f"Failed logging into victim {self.host} with user: {user}: {err}"
                 logger.debug(error_message)
-                self._publish_exploitation_event(False, error_message=error_message)
+                self._publish_exploitation_event(stamp, False, error_message=error_message)
                 self.report_login_attempt(False, user, current_password)
                 ssh.close()
                 continue
@@ -232,6 +237,7 @@ class SSHExploiter(HostExploiter):
         if status == ScanStatus.SCANNED:
             raise FailedExploitationError(self.exploit_result.error_message)
 
+        stamp = time()
         try:
             cmdline = f"{monkey_path_on_victim} {MONKEY_ARG}"
             cmdline += build_monkey_commandline(self.servers, self.current_depth + 1)
@@ -246,12 +252,12 @@ class SSHExploiter(HostExploiter):
             )
 
             self.exploit_result.propagation_success = True
-            self._publish_propagation_event(True)
+            self._publish_propagation_event(stamp, True)
             self.add_executed_cmd(cmdline)
 
         except Exception as exc:
             error_message = f"Error running monkey on victim {self.host}: ({exc})"
-            self._publish_exploitation_event(False, error_message=error_message)
+            self._publish_propagation_event(stamp, False, error_message=error_message)
             raise FailedExploitationError(error_message)
 
     def _is_port_open(self, ip: IPv4Address, port: int) -> bool:
@@ -311,6 +317,7 @@ class SSHExploiter(HostExploiter):
         monkey_path_on_victim: PurePath,
     ) -> ScanStatus:
         try:
+            stamp = time()
             with ssh.open_sftp() as ftp:
                 ftp.putfo(
                     agent_binary_file_object,
@@ -323,7 +330,7 @@ class SSHExploiter(HostExploiter):
             return ScanStatus.USED
         except Exception as exc:
             error_message = f"Error uploading file into victim {self.host}: ({exc})"
-            self._publish_propagation_event(False, error_message=error_message)
+            self._publish_propagation_event(stamp, False, error_message=error_message)
             self.exploit_result.error_message = error_message
             return ScanStatus.SCANNED