Agent: Fix twisted import parallelization bug

monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py

@@ -6,14 +6,16 @@ import threading
 import time
 from pathlib import Path
 
-from ldaptor.interfaces import IConnectedLDAPEntry
-from ldaptor.ldiftree import LDIFTreeEntry
 from ldaptor.protocols.ldap.ldapserver import LDAPServer
-from twisted.application import service
-from twisted.internet import reactor
 from twisted.internet.protocol import ServerFactory
-from twisted.python import log
+
-from twisted.python.components import registerAdapter
+# WARNING: It was observed that this LDAP server would raise an exception and fail to start if
+#          multiple Python threads attempt to start multiple LDAP servers simultaneously. It was
+#          thought that since each LDAP server is started in its own process, there would be no
+#          issue, however this is not the case. It seems that there may be something that is not
+#          thread- or multiprocess-safe about some of the twisted imports. Moving the twisted
+#          imports down into the functions where they are required and removing them from the top of
+#          this file appears to resolve the issue.
 
 logger = logging.getLogger(__name__)
 
@@ -32,6 +34,8 @@ class Tree:
     """
 
     def __init__(self, http_server_ip: str, http_server_port: int, storage_dir: Path):
+        from ldaptor.ldiftree import LDIFTreeEntry
+
         self.path = tempfile.mkdtemp(prefix="log4shell", suffix=".ldap", dir=storage_dir)
         self.db = LDIFTreeEntry(self.path)
 
@@ -91,14 +95,7 @@ class LDAPExploitServer:
         self._http_server_ip = http_server_ip
         self._http_server_port = http_server_port
         self._storage_dir = storage_dir
-
+        self._server_process = None
-        # A Twisted reactor can only be started and stopped once. It cannot be restarted after it
-        # has been stopped. To work around this, the reactor is configured and run in a separate
-        # process. This allows us to run multiple LDAP servers sequentially or simultaneously and
-        # stop each one when we're done with it.
-        self._server_process = multiprocessing.Process(
-            target=self._run_twisted_reactor, daemon=True
-        )
 
     def run(self):
         """
@@ -108,6 +105,15 @@ class LDAPExploitServer:
         :raises LDAPServerStartError: Indicates there was a problem starting the LDAP server.
         """
         logger.info("Starting LDAP exploit server")
+
+        # A Twisted reactor can only be started and stopped once. It cannot be restarted after it
+        # has been stopped. To work around this, the reactor is configured and run in a separate
+        # process. This allows us to run multiple LDAP servers sequentially or simultaneously and
+        # stop each one when we're done with it.
+        self._server_process = multiprocessing.Process(
+            target=self._run_twisted_reactor, daemon=True
+        )
+
         self._server_process.start()
         reactor_running = self._reactor_startup_completed.wait(REACTOR_START_TIMEOUT_SEC)
 
@@ -117,6 +123,8 @@ class LDAPExploitServer:
         logger.debug("The LDAP exploit server has successfully started")
 
     def _run_twisted_reactor(self):
+        from twisted.internet import reactor
+
         logger.debug(f"Starting log4shell LDAP server on port {self._ldap_server_port}")
         self._configure_twisted_reactor()
 
@@ -128,6 +136,8 @@ class LDAPExploitServer:
         reactor.run()
 
     def _check_if_reactor_startup_completed(self):
+        from twisted.internet import reactor
+
         check_interval_sec = 0.25
         num_checks = math.ceil(REACTOR_START_TIMEOUT_SEC / check_interval_sec)
 
@@ -141,6 +151,11 @@ class LDAPExploitServer:
             time.sleep(check_interval_sec)
 
     def _configure_twisted_reactor(self):
+        from ldaptor.interfaces import IConnectedLDAPEntry
+        from twisted.application import service
+        from twisted.internet import reactor
+        from twisted.python.components import registerAdapter
+
         LDAPExploitServer._output_twisted_logs_to_python_logger()
 
         registerAdapter(lambda x: x.root, LDAPServerFactory, IConnectedLDAPEntry)
@@ -155,6 +170,8 @@ class LDAPExploitServer:
 
     @staticmethod
     def _output_twisted_logs_to_python_logger():
+        from twisted.python import log
+
         # Configures Twisted to output its logs using the standard python logging module instead of
         # the Twisted logging module.
         # https://twistedmatrix.com/documents/current/api/twisted.python.log.PythonLoggingObserver.html