diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 425234760..b3e2c1d4c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -121,6 +121,11 @@ class InfectionMonkey: # TODO: `address_to_port()` should return the port as an integer. self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_port = int(self._cmd_island_port) + + self._island_address = SocketAddress( + IPv4Address(self._cmd_island_ip), self._cmd_island_port + ) + self._control_client = ControlClient( server_address=server, island_api_client=self._island_api_client ) @@ -231,10 +236,9 @@ class InfectionMonkey: config = self._control_channel.get_config() relay_port = get_free_tcp_port() - island_address = SocketAddress(IPv4Address(self._cmd_island_ip), self._cmd_island_port) self._relay = TCPRelay( relay_port, - island_address, + self._island_address, client_disconnect_timeout=config.keep_tunnel_open_time, ) @@ -488,7 +492,7 @@ class InfectionMonkey: def _close_tunnel(self): logger.info(f"Quitting tunnel {self._cmd_island_ip}") - notify_disconnect(self._cmd_island_ip, self._cmd_island_port) + notify_disconnect(self._island_address) def _send_log(self): monkey_log_path = get_agent_log_path() diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 3f7cc4629..f09a076f3 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -92,22 +92,26 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): def _send_remove_from_waitlist_control_message_to_relay(server: str): ip, port = address_to_ip_port(server) server_address = SocketAddress(IPv4Address(ip), int(port)) - notify_disconnect(IPv4Address(ip), int(port)) + notify_disconnect(server_address) -def notify_disconnect(server_ip: IPv4Address, server_port: int): +def notify_disconnect(server_address: SocketAddress): """ - Tell upstream relay that we no longer need the relay. + Tell upstream relay that we no longer need the relay - :param server_ip: The IP address of the server to notify. - :param server_port: The port of the server to notify. + :param server_address: The address of the server to notify """ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: d_socket.settimeout(LONG_REQUEST_TIMEOUT) try: - d_socket.connect((str(server_ip), server_port)) + d_socket.connect((str(server_address.ip), server_address.port)) d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) - logger.info(f"Control message was sent to the server/relay {server_ip}:{server_port}") + logger.info( + "Control message was sent to the server/relay " + f"{server_address.ip}:{server_address.port}" + ) except OSError as err: - logger.error(f"Error connecting to socket {server_ip}:{server_port}: {err}") + logger.error( + f"Error connecting to socket {server_address.ip}:{server_address.port}: {err}" + )