From 76fe0be990317bc91b7839c7115f8eeb950d5d4b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 14:54:13 -0400 Subject: [PATCH] Agent: Address addr type TODO in RelayConnectionHandler --- .../network/relay/relay_connection_handler.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network/relay/relay_connection_handler.py b/monkey/infection_monkey/network/relay/relay_connection_handler.py index fe08a5a61..156b78ddd 100644 --- a/monkey/infection_monkey/network/relay/relay_connection_handler.py +++ b/monkey/infection_monkey/network/relay/relay_connection_handler.py @@ -13,10 +13,13 @@ class RelayConnectionHandler: self._relay_user_handler = relay_user_handler def handle_new_connection(self, sock: socket.socket): + addr, _ = sock.getpeername() + addr = IPv4Address(addr) + control_message = sock.recv(socket.MSG_PEEK) - addr, _ = sock.getpeername() # TODO check the type of the addr object + if control_message.startswith(RELAY_CONTROL_MESSAGE): - self._relay_user_handler.disconnect_user(IPv4Address(addr)) + self._relay_user_handler.disconnect_user(addr) else: - self._relay_user_handler.add_relay_user(IPv4Address(addr)) + self._relay_user_handler.add_relay_user(addr) self._pipe_spawner.spawn_pipe(sock)