Agent: Clean up RelayUserHandler

Add get_potential_users and remove pipe handling

monkey/infection_monkey/network/relay/relay_user_handler.py

@@ -1,10 +1,8 @@
-import socket
 from dataclasses import dataclass
 from ipaddress import IPv4Address
+from threading import Lock
 from time import time
-from typing import Callable, Dict
-
-from threding import Lock
+from typing import Dict
 
 RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -"
 
@@ -16,19 +14,18 @@ class RelayUser:
 
 
 class RelayUserHandler:
-    def __init__(self, spawn_new_pipe: Callable[[socket.socket, IPv4Address], None]):
+    def __init__(self):
         self._relay_users: Dict[IPv4Address, RelayUser] = {}
         self._potential_users: Dict[IPv4Address, RelayUser] = {}
-        self._spawn_new_pipe = spawn_new_pipe
 
         self._lock = Lock()
 
-    def add_relay_user(self, source_socket: socket.socket, user_address: IPv4Address):
+    def add_relay_user(self, user_address: IPv4Address):
         """
         Handle new user connection.
 
         :param source_socket: A source socket
-        :param user_addres: An address defining RelayUser which will be added to the relay
+        :param user_address: An address defining RelayUser which will be added to the relay
         """
 
         with self._lock:
@@ -36,7 +33,6 @@ class RelayUserHandler:
                 del self._potential_users[user_address]
 
             self._relay_users[user_address] = RelayUser(user_address, time())
-            self._spawn_new_pipe(source_socket, user_address)
 
     def add_potential_user(self, user_address: IPv4Address):
         """
@@ -48,12 +44,12 @@ class RelayUserHandler:
         with self._lock:
             self._potential_users[user_address] = RelayUser(user_address, time())
 
-    def on_user_data_recieved(self, data: bytes, user_address: IPv4Address) -> bool:
+    def on_user_data_received(self, data: bytes, user_address: IPv4Address) -> bool:
         """
         Disconnect a user with a specific starting data.
 
-        :param data: The data that a relay recieved
-        :param user_address: An address defining RelayUser which recieved the data
+        :param data: The data that a relay received
+        :param user_address: An address defining RelayUser which received the data
         """
         if data.startswith(RELAY_CONTROL_MESSAGE):
             self._disconnect_user(user_address)
@@ -64,3 +60,7 @@ class RelayUserHandler:
         with self._lock:
             if user_address in self._relay_users:
                 del self._relay_users[user_address]
+
+    def get_potential_users(self) -> Dict[IPv4Address, RelayUser]:
+        with self._lock:
+            return self._potential_users.copy()