From 5c7c54e45045234330aa06111e642c0a359e7b97 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 7 Sep 2022 12:17:49 +0000 Subject: [PATCH] Agent: Use timeout in build_tcprelay_deps() --- monkey/infection_monkey/network/relay/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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(