diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py index a3a83d9a2..41bc81e80 100644 --- a/monkey/infection_monkey/network/relay/__init__.py +++ b/monkey/infection_monkey/network/relay/__init__.py @@ -1,3 +1,4 @@ from .relay_user_handler import RelayUser, RelayUserHandler from .tcp_connection_handler import TCPConnectionHandler from .tcp import SocketsPipe +from .tcp_pipe_spawner import TCPPipeSpawner diff --git a/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py b/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py new file mode 100644 index 000000000..16523dfe3 --- /dev/null +++ b/monkey/infection_monkey/network/relay/tcp_pipe_spawner.py @@ -0,0 +1,25 @@ +import socket +from ipaddress import IPv4Address +from typing import Callable + +from .tcp import SocketsPipe + + +class TCPPipeSpawner: + def __init__(self, target_addr: IPv4Address, target_port: int): + self._target_addr = target_addr + self._target_port = target_port + + def spawn_pipe(self, source: socket.socket) -> SocketsPipe: + dest = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + dest.connect((self._target_addr, self._target_port)) + except socket.error as err: + source.close() + dest.close() + raise err + + return SocketsPipe(source, dest, client_data_received=self._client_data_received) + + def notify_client_data_received(self, callback: Callable[[bytes], bool]): + self._client_data_received = callback