diff --git a/monkey/common/network/network_utils.py b/monkey/common/network/network_utils.py index f686268e0..13a54dce1 100644 --- a/monkey/common/network/network_utils.py +++ b/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 diff --git a/monkey/infection_monkey/network/tools.py b/monkey/infection_monkey/network/tools.py index 9bf69e28a..1f0c276e3 100644 --- a/monkey/infection_monkey/network/tools.py +++ b/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]: