Agent: Add TCPConnectionHandler

monkey/infection_monkey/network/relay/tcp_connection_handler.py

@@ -0,0 +1,43 @@
+import socket
+from ipaddress import IPv4Address
+from threading import Thread
+from typing import Callable
+
+PROXY_TIMEOUT = 2.5
+
+
+class TCPConnectionHandler(Thread):
+    """Accepts connections on a TCP socket."""
+
+    def __init__(
+        self,
+        local_port: int,
+        local_host: str = "",
+        client_connected: Callable[[socket.socket, IPv4Address], None] = None,
+    ):
+        self.local_port = local_port
+        self.local_host = local_host
+        self._client_connected = client_connected
+        super().__init__()
+        self.daemon = True
+        self._stopped = False
+
+    def run(self):
+        l_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        l_socket.bind((self.local_host, self.local_port))
+        l_socket.settimeout(PROXY_TIMEOUT)
+        l_socket.listen(5)
+
+        while not self._stopped:
+            try:
+                source, address = l_socket.accept()
+            except socket.timeout:
+                continue
+
+            if self._client_connected:
+                self._client_connected(source, IPv4Address(address[0]))
+
+        l_socket.close()
+
+    def stop(self):
+        self._stopped = True