diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py index ad5d325c5..61728b643 100644 --- a/monkey/infection_monkey/network/relay/__init__.py +++ b/monkey/infection_monkey/network/relay/__init__.py @@ -1,4 +1,7 @@ -from .relay_connection_handler import RelayConnectionHandler, RELAY_CONTROL_MESSAGE +from .relay_connection_handler import ( + RelayConnectionHandler, + RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST, +) from .relay_user_handler import RelayUser, RelayUserHandler from .sockets_pipe import SocketsPipe from .tcp_connection_handler import TCPConnectionHandler diff --git a/monkey/infection_monkey/network/relay/relay_connection_handler.py b/monkey/infection_monkey/network/relay/relay_connection_handler.py index 3d7755fb9..4b4475e52 100644 --- a/monkey/infection_monkey/network/relay/relay_connection_handler.py +++ b/monkey/infection_monkey/network/relay/relay_connection_handler.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address from .relay_user_handler import RelayUserHandler from .tcp_pipe_spawner import TCPPipeSpawner -RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -" +RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -" class RelayConnectionHandler: @@ -25,7 +25,7 @@ class RelayConnectionHandler: control_message = sock.recv(socket.MSG_PEEK) - if control_message.startswith(RELAY_CONTROL_MESSAGE): + if control_message.startswith(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST): self._relay_user_handler.disconnect_user(addr) else: self._relay_user_handler.add_relay_user(addr) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 333eafd6d..096500c5b 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -6,7 +6,7 @@ import requests from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import address_to_ip_port -from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE +from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) @@ -57,6 +57,6 @@ def _send_remove_from_waitlist_control_message_to_relay(server: str): try: d_socket.connect((ip, int(port))) - d_socket.send(RELAY_CONTROL_MESSAGE) + d_socket.send(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) except OSError as err: logger.error(f"Error connecting to socket {server}: {err}")