From a0f566ef49b4b477c2e605f810de1312fef38484 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 31 Aug 2022 18:44:18 +0000
Subject: [PATCH] Agent: Add network utilities for connections

---
 monkey/infection_monkey/network/tools.py      | 37 +++++++++++++++++++
 .../network/test_network_tools.py             | 13 +++++++
 2 files changed, 50 insertions(+)
 create mode 100644 monkey/tests/unit_tests/infection_monkey/network/test_network_tools.py

diff --git a/monkey/infection_monkey/network/tools.py b/monkey/infection_monkey/network/tools.py
index c612a7e48..9bf69e28a 100644
--- a/monkey/infection_monkey/network/tools.py
+++ b/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
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
new file mode 100644
index 000000000..77a4e8b00
--- /dev/null
+++ b/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"])