Agent: Add class documentation to ExploitClassHTTPServer

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()