From 1f723b174eb9c91c186921d4d9cc35a8491ac0a8 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Thu, 1 Sep 2022 13:54:17 +0000
Subject: [PATCH] Agent: Add TCPConnectionHandler

---
 .../network/relay/tcp_connection_handler.py   | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 monkey/infection_monkey/network/relay/tcp_connection_handler.py

diff --git a/monkey/infection_monkey/network/relay/tcp_connection_handler.py b/monkey/infection_monkey/network/relay/tcp_connection_handler.py
new file mode 100644
index 000000000..9b0c4f023
--- /dev/null
+++ b/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