Agent: Create HTTP handler class dynamically for ExploitClassHTTPServer

monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py

@@ -7,31 +7,16 @@ logger = logging.getLogger(__name__)
 HTTP_TOO_MANY_REQUESTS_ERROR_CODE = 429
 
 
-# If we need to run multiple HTTP servers in parallel, we'll need to either:
-#   1. Use multiprocessing so that each HTTPHandler class has its own class_downloaded variable
-#   2. Create a metaclass and define the handler class dymanically at runtime
-class HTTPHandler(http.server.BaseHTTPRequestHandler):
-
-    java_class: bytes
-    class_downloaded: threading.Event
-    download_lock: threading.Lock
-
-    @classmethod
-    def initialize(cls, java_class: bytes, class_downloaded: threading.Event):
-        cls.java_class = java_class
-        cls.class_downloaded = class_downloaded
-        cls.download_lock = threading.Lock()
-
-    def do_GET(self):
-        with HTTPHandler.download_lock:
-            if HTTPHandler.class_downloaded.is_set():
+def do_GET(self):
+    with self.download_lock:
+        if self.class_downloaded.is_set():
             self.send_error(
                 HTTP_TOO_MANY_REQUESTS_ERROR_CODE,
                 "Java exploit class has already been downloaded",
             )
             return
 
-            HTTPHandler.class_downloaded.set()
+        self.class_downloaded.set()
 
     logger.info("Java class server received a GET request!")
     self.send_response(200)
@@ -41,6 +26,19 @@ class HTTPHandler(http.server.BaseHTTPRequestHandler):
     self.wfile.write(self.java_class)
 
 
+def get_new_http_handler_class(java_class: bytes, class_downloaded: threading.Event):
+    return type(
+        "http_handler_class",
+        (http.server.BaseHTTPRequestHandler,),
+        {
+            "java_class": java_class,
+            "class_downloaded": class_downloaded,
+            "download_lock": threading.Lock(),
+            "do_GET": do_GET,
+        },
+    )
+
+
 class ExploitClassHTTPServer:
     """
     An HTTP server that serves Java bytecode for use with the Log4Shell exploiter. This server
@@ -62,9 +60,9 @@ class ExploitClassHTTPServer:
         self._class_downloaded = threading.Event()
         self._poll_interval = poll_interval
 
-        HTTPHandler.initialize(java_class, self._class_downloaded)
+        http_handler_class = get_new_http_handler_class(java_class, self._class_downloaded)
 
-        self._server = http.server.HTTPServer((ip, port), HTTPHandler)
+        self._server = http.server.HTTPServer((ip, port), http_handler_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()`