diff --git a/monkey/infection_monkey/tcp_relay.py b/monkey/infection_monkey/tcp_relay.py index 9fa4be1b6..243204b76 100644 --- a/monkey/infection_monkey/tcp_relay.py +++ b/monkey/infection_monkey/tcp_relay.py @@ -6,6 +6,7 @@ from typing import List from infection_monkey.transport.tcp import TcpProxy DEFAULT_NEW_CLIENT_TIMEOUT = 3 # Wait up to 3 seconds for potential new clients to connect +RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -" @dataclass @@ -77,7 +78,7 @@ class TCPRelay(Thread): self._potential_users.append(RelayUser(user, time())) def on_user_data_received(self, data: bytes, user: str) -> bool: - if data.startswith(b"-"): + if data.startswith(RELAY_CONTROL_MESSAGE): self._disconnect_user(user) return False return True diff --git a/monkey/tests/unit_tests/infection_monkey/test_tcp_relay.py b/monkey/tests/unit_tests/infection_monkey/test_tcp_relay.py index 48dcce2a6..02c577797 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_tcp_relay.py +++ b/monkey/tests/unit_tests/infection_monkey/test_tcp_relay.py @@ -1,6 +1,6 @@ from threading import Thread -from monkey.infection_monkey.tcp_relay import TCPRelay +from monkey.infection_monkey.tcp_relay import RELAY_CONTROL_MESSAGE, TCPRelay def join_or_kill_thread(thread: Thread, timeout: float): @@ -46,7 +46,7 @@ def test_user_removed_on_request(): relay = TCPRelay(9975, "0.0.0.0", 9976) new_user = "0.0.0.1" relay.on_user_connected(new_user) - relay.on_user_data_received(b"-", "0.0.0.1") + relay.on_user_data_received(RELAY_CONTROL_MESSAGE, "0.0.0.1") users = relay.relay_users() assert len(users) == 0