diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 728c463a0..d2ea3db0b 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -43,7 +43,6 @@ from infection_monkey.model import VictimHostFactory from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_free_tcp_port, get_network_interfaces from infection_monkey.network.relay import TCPRelay -from infection_monkey.network.tools import connect from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter @@ -183,12 +182,10 @@ class InfectionMonkey: config = control_channel.get_config() local_port = get_free_tcp_port() - sock, ip_str, port = connect(self._opts.servers) - sock.close() self._relay = TCPRelay( local_port, - IPv4Address(ip_str), - port, + IPv4Address(self._cmd_island_ip), + self._cmd_island_port, client_disconnect_timeout=config.keep_tunnel_open_time, ) diff --git a/monkey/infection_monkey/network/tools.py b/monkey/infection_monkey/network/tools.py index 1f0c276e3..c612a7e48 100644 --- a/monkey/infection_monkey/network/tools.py +++ b/monkey/infection_monkey/network/tools.py @@ -3,10 +3,8 @@ import select import socket import struct import sys -from typing import List, Tuple, Union from common.common_consts.timeouts import CONNECTION_TIMEOUT -from common.network.network_utils import address_to_ip_port from infection_monkey.network.info import get_routes DEFAULT_TIMEOUT = CONNECTION_TIMEOUT @@ -92,40 +90,3 @@ def get_interface_to_target(dst): paths.sort() ret = paths[-1][1] return ret[1] - - -def connect(connections: List[str]) -> Tuple[socket.socket, str, int]: - """ - Attempt to connect to addresses in the given list. - - :param connections: The addresses to try and connect to. - :return: The socket, address, and port of the connection. - :raises: ConnectionError if no connection could be established. - :raises: ValueError if an improper connection is provided. - """ - for connection in connections: - ip, port = address_to_ip_port(connection) - if port is None: - raise ValueError("Connection does not contain a port") - sock = try_connect(ip, int(port)) - if sock: - return sock, ip, int(port) - - raise ConnectionError("Could not connect to a server in the server list") - - -def try_connect(ip: str, port: int) -> Union[socket.socket, None]: - """ - Attempt to establish a connection. - - :param ip: The IP to use. - :param port: The port to use. - :return: The socket on a successful connection, otherwise None. - """ - try: - logging.debug(f"Attempting to connect to {ip}:{port}") - sock = socket.create_connection((ip, port), timeout=1) - except Exception: - return None - - return sock diff --git a/monkey/tests/unit_tests/infection_monkey/network/test_network_tools.py b/monkey/tests/unit_tests/infection_monkey/network/test_network_tools.py deleted file mode 100644 index 77a4e8b00..000000000 --- a/monkey/tests/unit_tests/infection_monkey/network/test_network_tools.py +++ /dev/null @@ -1,13 +0,0 @@ -import pytest - -from infection_monkey.network.tools import connect - - -def test_connect_raises_with_empty_list(): - with pytest.raises(ConnectionError): - connect([]) - - -def test_connect_raises_with_bad_data(): - with pytest.raises(ValueError): - connect(["no-port"])