diff --git a/monkey/common/network/network_utils.py b/monkey/common/network/network_utils.py index 13a54dce1..d211474a5 100644 --- a/monkey/common/network/network_utils.py +++ b/monkey/common/network/network_utils.py @@ -1,6 +1,7 @@ from typing import Optional, Tuple +# TODO: `address_to_port()` should return the port as an integer. def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]: """ Split a string containing an IP address (and optionally a port) into IP and Port components. diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index c96b2f159..005fad62e 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -103,7 +103,9 @@ class InfectionMonkey: # TODO: Revisit variable names server = self._get_server() + # 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_ip = int(self._cmd_island_ip) self._control_client = ControlClient(server_address=server) # TODO Refactor the telemetry messengers to accept control client @@ -189,7 +191,7 @@ class InfectionMonkey: self._relay = TCPRelay( relay_port, IPv4Address(self._cmd_island_ip), - int(self._cmd_island_port), + self._cmd_island_port, client_disconnect_timeout=config.keep_tunnel_open_time, ) @@ -427,7 +429,7 @@ class InfectionMonkey: def _close_tunnel(self): logger.info(f"Quitting tunnel {self._cmd_island_ip}") - notify_disconnect(self._cmd_island_ip, int(self._cmd_island_port)) + notify_disconnect(self._cmd_island_ip, self._cmd_island_port) def _send_log(self): monkey_log_path = get_agent_log_path()