diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 3db3cb856..425234760 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,6 +21,7 @@ from common.network.network_utils import ( get_my_ip_addresses, get_network_interfaces, ) +from common.types import SocketAddress from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version @@ -230,10 +231,10 @@ 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, - IPv4Address(self._cmd_island_ip), - self._cmd_island_port, + island_address, client_disconnect_timeout=config.keep_tunnel_open_time, ) diff --git a/monkey/infection_monkey/network/relay/tcp_relay.py b/monkey/infection_monkey/network/relay/tcp_relay.py index c2663cb1d..b25600752 100644 --- a/monkey/infection_monkey/network/relay/tcp_relay.py +++ b/monkey/infection_monkey/network/relay/tcp_relay.py @@ -3,6 +3,7 @@ from logging import getLogger from threading import Lock, Thread from time import sleep +from common.types import SocketAddress from infection_monkey.network.relay import ( RelayConnectionHandler, RelayUserHandler, @@ -22,15 +23,14 @@ class TCPRelay(Thread, InterruptableThreadMixin): def __init__( self, relay_port: int, - dest_addr: IPv4Address, - dest_port: int, + dest_address: SocketAddress, client_disconnect_timeout: float, ): self._user_handler = RelayUserHandler( new_client_timeout=client_disconnect_timeout, client_disconnect_timeout=client_disconnect_timeout, ) - self._pipe_spawner = TCPPipeSpawner(dest_addr, dest_port) + self._pipe_spawner = TCPPipeSpawner(dest_address.ip, dest_address.port) relay_filter = RelayConnectionHandler(self._pipe_spawner, self._user_handler) self._connection_handler = TCPConnectionHandler( bind_host="",