Agent: Add network utilities for connections

monkey/infection_monkey/network/tools.py

@@ -3,6 +3,7 @@ import select
 import socket
 import struct
 import sys
+from typing import List, Tuple, Union
 
 from common.common_consts.timeouts import CONNECTION_TIMEOUT
 from infection_monkey.network.info import get_routes
@@ -90,3 +91,39 @@ 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 = connection.rpartition(":")
+        ip = ip.strip("[]")
+        sock = try_connect(ip, int(port))
+        if sock:
+            return sock, ip, int(port)
+
+    raise ConnectionError
+
+
+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

@@ -0,0 +1,13 @@
+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"])