diff --git a/monkey/infection_monkey/network/relay/relay_user_handler.py b/monkey/infection_monkey/network/relay/relay_user_handler.py index ba1855655..364f5985e 100644 --- a/monkey/infection_monkey/network/relay/relay_user_handler.py +++ b/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()