diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index e4321ad43..81bb00203 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,14 +1,14 @@ -from . import TCPConnectionHandler, TCPPipeSpawner, RelayUserHandler, RelayConnectionHandler from ipaddress import IPv4Address from typing import Tuple +from . import RelayConnectionHandler, RelayUserHandler, TCPConnectionHandler, TCPPipeSpawner + def build_tcprelay_deps( local_port: int, dest_addr: IPv4Address, dest_port: int, client_disconnect_timeout: float ) -> Tuple[RelayUserHandler, TCPPipeSpawner, TCPConnectionHandler]: - # TODO: Add the timeouts - relay_user_handler = RelayUserHandler() + relay_user_handler = RelayUserHandler(client_disconnect_timeout=client_disconnect_timeout) pipe_spawner = TCPPipeSpawner(dest_addr, dest_port) relay_filter = RelayConnectionHandler(pipe_spawner, relay_user_handler) connection_handler = TCPConnectionHandler(