Agent: Remove connect()

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),
+            IPv4Address(self._cmd_island_ip),
-            port,
+            self._cmd_island_port,
             client_disconnect_timeout=config.keep_tunnel_open_time,
         )
 

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

monkey/tests/unit_tests/infection_monkey/network/test_network_tools.py

@@ -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"])