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 c93f910e6..9b82de1d6 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 @@ -7,6 +7,9 @@ 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 @@ -29,7 +32,21 @@ class HTTPHandler(http.server.BaseHTTPRequestHandler): class ExploitClassHTTPServer: + """ + An HTTP server that serves Java bytecode for use with the Log4Shell exploiter. This server + limits the number of requests to one. That is, after one victim has downloaded the java + bytecode, the server will respond with a 429 error to all future requests. + + Note: There can only be one instance of this class at a time due to the way it is implemented. + """ + def __init__(self, ip: str, port: int, java_class: bytes, poll_interval: float = 0.5): + """ + :param ip: The IP address that the server will bind to + :param port: The port that the server will listen on + :param java_class: The compiled Java bytecode that the server will serve + :param poll_interval: Poll for shutdown every `poll_interval` seconds, defaults to 0.5. + """ logger.debug(f"The Java Exploit class will be served at {ip}:{port}") self._class_downloaded = threading.Event() @@ -51,17 +68,29 @@ class ExploitClassHTTPServer: HTTPHandler.class_downloaded = self._class_downloaded def run(self): + """ + Runs the HTTP server in the background and blocks until the server has started. + """ logger.info("Starting ExploitClassHTTPServer") self._class_downloaded.clear() # NOTE: Unlike in LDAPExploitServer, we theoretically don't need to worry about a race # between when `serve_forever()` is ready to handle requests and when the victim machine - # sends its requests. See + # sends its requests. This could change if we switch from multithreading to multiprocessing. + # See # https://stackoverflow.com/questions/22606480/how-can-i-test-if-python-http-server-httpserver-is-serving-forever # for more information. self._server_thread.start() def stop(self, timeout: float = None): + """ + Stops the HTTP server. + + :param timeout: A floating point number of seconds to wait for the server to stop. If this + argument is None (the default), the method blocks until the HTTP server + terminates. If `timeout` is a positive floating point number, this method + blocks for at most `timeout` seconds. + """ if self._server_thread.is_alive(): logger.debug("Stopping the Java Exploit class HTTP server") self._server.shutdown() @@ -73,4 +102,11 @@ class ExploitClassHTTPServer: logger.debug("The Java Exploit class HTTP server has stopped") def exploit_class_downloaded(self) -> bool: + """ + Returns whether or not a victim has downloaded the Java bytecode from the server. + + :return: True if the victim has downloaded the Java bytecode from the server. False + otherwise. + :rtype: bool + """ return self._class_downloaded.is_set()