From 6b934d6de56a6a538438fe1bced6e4f2419c8aea Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 13 Jan 2022 08:56:23 -0500
Subject: [PATCH] Agent: Wrap log4shell LDAP server in a process

 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.
---
 .../exploit/log4shell_utils/ldap_server.py    | 47 ++++++++++++++-----
 1 file changed, 34 insertions(+), 13 deletions(-)

diff --git a/monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py b/monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py
index f0acac141..382993d68 100644
--- a/monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py
+++ b/monkey/infection_monkey/exploit/log4shell_utils/ldap_server.py
@@ -1,3 +1,5 @@
+import logging
+import multiprocessing
 import tempfile
 from pathlib import Path
 
@@ -10,6 +12,8 @@ from twisted.internet.protocol import ServerFactory
 from twisted.python import log
 from twisted.python.components import registerAdapter
 
+logger = logging.getLogger(__name__)
+
 EXPLOIT_RDN = "dn=Exploit"
 
 
@@ -63,33 +67,50 @@ class LDAPExploitServer:
     def __init__(
         self, ldap_server_port: int, http_server_ip: str, http_server_port: int, storage_dir: Path
     ):
-        LDAPExploitServer.output_twisted_logs_to_python_logger()
+        self._ldap_server_port = ldap_server_port
+        self._http_server_ip = http_server_ip
+        self._http_server_port = http_server_port
+        self._storage_dir = storage_dir
+
+        # 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):
+        self._server_process.start()
+        self._server_process.join()
+
+    def _run_twisted_reactor(self):
+        self._configure_twisted_reactor()
+        logger.debug(f"Starting log4shell LDAP server on port {self._ldap_server_port}")
+        reactor.run()
+
+    def _configure_twisted_reactor(self):
+        LDAPExploitServer._output_twisted_logs_to_python_logger()
 
         registerAdapter(lambda x: x.root, LDAPServerFactory, IConnectedLDAPEntry)
 
-        tree = Tree(http_server_ip, http_server_port, storage_dir)
+        tree = Tree(self._http_server_ip, self._http_server_port, self._storage_dir)
         factory = LDAPServerFactory(tree.db)
         factory.debug = True
 
         application = service.Application("ldaptor-server")
         service.IServiceCollection(application)
-        reactor.listenTCP(ldap_server_port, factory)
+        reactor.listenTCP(self._ldap_server_port, factory)
 
     @staticmethod
-    def output_twisted_logs_to_python_logger():
+    def _output_twisted_logs_to_python_logger():
         # 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
         log_observer = log.PythonLoggingObserver()
         log_observer.start()
 
-    def run(self):
-        # For an explaination of installSignalHandlers=0, see
-        # https://twistedmatrix.com/trac/wiki/FrequentlyAskedQuestions#Igetexceptions.ValueError:signalonlyworksinmainthreadwhenItrytorunmyTwistedprogramWhatswrong
-        reactor.run(installSignalHandlers=0)
-
     def stop(self):
-        # Since `reactor.run()` may not be running on the main thread, reactor.stop() must be
-        # invoked with reactor.callFromThread()
-        # https://twistedmatrix.com/documents/12.2.0/core/howto/threading.html
-        reactor.callFromThread(reactor.stop)
+        # The Twisted reactor registers signal handlers so it can catch SIGTERM and gracefully
+        # shutdown.
+        self._server_process.terminate()