Agent: Use addres_to_ip_port in connect()

monkey/common/network/network_utils.py

@@ -2,6 +2,13 @@ from typing import Optional, Tuple
 
 
 def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]:
+    """
+    Split a string containing an IP address (and optionally a port) into IP and Port components.
+    Currently only works for IPv4 addresses.
+
+    :param address: The address string.
+    :return: Tuple of IP and port strings. The port may be None if no port was in the address.
+    """
     if ":" in address:
         ip, port = address.split(":")
         return ip, port or None

monkey/infection_monkey/network/tools.py

@@ -6,6 +6,7 @@ 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
@@ -103,13 +104,14 @@ def connect(connections: List[str]) -> Tuple[socket.socket, str, int]:
     :raises: ValueError if an improper connection is provided.
     """
     for connection in connections:
-        ip, _, port = connection.rpartition(":")
-        ip = ip.strip("[]")
+        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
+    raise ConnectionError("Could not connect to a server in the server list")
 
 
 def try_connect(ip: str, port: int) -> Union[socket.socket, None]: