diff --git a/monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py b/monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py index 3e0cdd38d..89eee7808 100644 --- a/monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py +++ b/monkey/infection_monkey/exploit/log4shell_utils/exploit_class_http_server.py @@ -1,6 +1,7 @@ import http.server import logging import threading +from typing import Type logger = logging.getLogger(__name__) @@ -26,7 +27,18 @@ def do_GET(self): self.wfile.write(self.java_class) -def get_new_http_handler_class(java_class: bytes, class_downloaded: threading.Event): +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,), @@ -60,7 +72,7 @@ class ExploitClassHTTPServer: self._class_downloaded = threading.Event() self._poll_interval = poll_interval - HTTPHandler = get_new_http_handler_class(java_class, self._class_downloaded) + HTTPHandler = _get_new_http_handler_class(java_class, self._class_downloaded) self._server = http.server.HTTPServer((ip, port), HTTPHandler) # Setting `daemon=True` to save ourselves some trouble when this is merged to the