From 016bf5c795e3e00c9a711abd20348d41ee9d4c67 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 5 Oct 2022 19:55:56 +0000
Subject: [PATCH] Agent: Stamp times before the exploit runs

---
 monkey/infection_monkey/exploit/log4shell.py | 54 ++++++++------------
 1 file changed, 21 insertions(+), 33 deletions(-)

diff --git a/monkey/infection_monkey/exploit/log4shell.py b/monkey/infection_monkey/exploit/log4shell.py
index 9b4aec6f5..c5bee65a8 100644
--- a/monkey/infection_monkey/exploit/log4shell.py
+++ b/monkey/infection_monkey/exploit/log4shell.py
@@ -1,6 +1,7 @@
 import logging
 import time
 from pathlib import PurePath
+from typing import Tuple
 
 from common import OperatingSystem
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT
@@ -28,7 +29,6 @@ from infection_monkey.utils.threading import interruptible_iter
 logger = logging.getLogger(__name__)
 
 LOG4SHELL_EXPLOITER_TAG = "log4shell-exploiter"
-PROPAGATION_TAGS = (LOG4SHELL_EXPLOITER_TAG, T1203_ATTACK_TECHNIQUE_TAG, T1105_ATTACK_TECHNIQUE_TAG)
 
 
 class Log4ShellExploiter(WebRCE):
@@ -36,6 +36,12 @@ class Log4ShellExploiter(WebRCE):
     SERVER_SHUTDOWN_TIMEOUT = LONG_REQUEST_TIMEOUT
     REQUEST_TO_VICTIM_TIMEOUT = MEDIUM_REQUEST_TIMEOUT
 
+    def _exploiter_tags(self) -> Tuple[str, ...]:
+        return (LOG4SHELL_EXPLOITER_TAG, T1203_ATTACK_TECHNIQUE_TAG)
+
+    def _propagation_tags(self) -> Tuple[str, ...]:
+        return (LOG4SHELL_EXPLOITER_TAG, T1203_ATTACK_TECHNIQUE_TAG, T1105_ATTACK_TECHNIQUE_TAG)
+
     def _exploit_host(self) -> ExploiterResultData:
         self._open_ports = [
             int(port[0]) for port in WebRCE.get_open_service_ports(self.host, self.HTTP, ["http"])
@@ -43,10 +49,6 @@ class Log4ShellExploiter(WebRCE):
 
         if not self._open_ports:
             logger.info("Could not find any open web ports to exploit")
-            self._publish_exploitation_event(
-                target=self.host.ip_addr,
-                exploitation_success=False,
-            )
             return self.exploit_result
 
         self._configure_servers()
@@ -156,31 +158,34 @@ class Log4ShellExploiter(WebRCE):
                     f"on port {port}"
                 )
                 try:
+                    timestamp = time.time()
                     url = exploit.trigger_exploit(self._build_ldap_payload(), self.host, port)
                 except Exception as err:
-                    error_message = "An error occurred while attempting to exploit log4shell on a "
-                    f"potential {exploit.service_name} service: {err}"
+                    error_message = (
+                        "An error occurred while attempting to exploit log4shell on a "
+                        f"potential {exploit.service_name} service: {err}"
+                    )
 
                     logger.warning(error_message)
 
-                    self._publish_exploitation_event(
-                        target=self.host.ip_addr,
-                        exploitation_success=False,
-                        error_message=error_message,
-                        tags=(LOG4SHELL_EXPLOITER_TAG,),
-                    )
+                    self._publish_exploitation_event(timestamp, False, error_message=error_message)
 
-                if self._wait_for_victim():
+                if self._wait_for_victim(timestamp):
                     self.exploit_info["vulnerable_service"] = {
                         "service_name": exploit.service_name,
                         "port": port,
                     }
                     self.exploit_info["vulnerable_urls"].append(url)
 
-    def _wait_for_victim(self) -> bool:
+    def _wait_for_victim(self, timestamp: float) -> bool:
         victim_called_back = self._wait_for_victim_to_download_java_bytecode()
         if victim_called_back:
+            self._publish_exploitation_event(timestamp, True)
             self._wait_for_victim_to_download_agent()
+        else:
+            error_message = "Timed out while waiting for victim to download the java bytecode"
+            logger.debug(error_message)
+            self._publish_exploitation_event(timestamp, False, error_message=error_message)
 
         return victim_called_back
 
@@ -190,24 +195,11 @@ class Log4ShellExploiter(WebRCE):
 
         while not timer.is_expired():
             if self._exploit_class_http_server.exploit_class_downloaded():
-                self._publish_exploitation_event(
-                    target=self.host.ip_addr,
-                    exploitation_success=True,
-                    tags=(LOG4SHELL_EXPLOITER_TAG, T1203_ATTACK_TECHNIQUE_TAG),
-                )
                 self.exploit_result.exploitation_success = True
                 return True
 
             time.sleep(1)
 
-        error_message = "Timed out while waiting for victim to download the java bytecode"
-        logger.debug(error_message)
-        self._publish_exploitation_event(
-            target=self.host.ip_addr,
-            exploitation_success=False,
-            error_message=error_message,
-            tags=(LOG4SHELL_EXPLOITER_TAG, T1203_ATTACK_TECHNIQUE_TAG),
-        )
         return False
 
     def _wait_for_victim_to_download_agent(self):
@@ -216,11 +208,7 @@ class Log4ShellExploiter(WebRCE):
 
         while not timer.is_expired():
             if self._agent_http_server_thread.downloads > 0:
-                self._publish_propagation_event(
-                    target=self.host.ip_addr,
-                    propagation_success=True,
-                    tags=PROPAGATION_TAGS,
-                )
+                self._publish_propagation_event(success=True)
                 self.exploit_result.propagation_success = True
                 break