add test_connection_util

simulation-system/libs/csle-common/tests/test_connection_util.py

@@ -0,0 +1,319 @@
+from unittest.mock import patch, MagicMock
+from csle_common.util.connection_util import ConnectionUtil
+from csle_common.dao.emulation_observation.common.emulation_connection_observation_state import (
+    EmulationConnectionObservationState,
+)
+import csle_common.constants.constants as constants
+
+
+class TestConnectionUtilSuite:
+    """
+    Test suite for connection_util
+    """
+
+    @patch("paramiko.SSHClient")
+    def test_ssh_setup_connection(self, mock_SSHClient) -> None:
+        """
+        Test the helper function for setting up a SSH connection
+
+        :param mock_SSHClient: mock_SSHClient
+
+        :return: None
+        """
+        mock_ssh_client = MagicMock()
+        mock_SSHClient.return_value = mock_ssh_client
+        mock_transport = MagicMock()
+        mock_relay_channel = MagicMock()
+        mock_ssh_client.get_transport.return_value = mock_transport
+        mock_transport.open_channel.return_value = mock_relay_channel
+
+        a = MagicMock()
+        a.ips = ["192.168.1.2"]
+        a.ips_match.side_effect = lambda ips: True
+
+        credentials = [MagicMock()]
+        credentials[0].service = constants.SSH.SERVICE_NAME
+        credentials[0].username = "user"
+        credentials[0].pw = "password"
+        credentials[0].port = 22
+
+        proxy_connections = [MagicMock()]
+        proxy_connections[0].ip = "192.168.1.1"
+        proxy_connections[0].conn.get_transport.return_value = mock_transport
+
+        s = MagicMock()
+        s.emulation_env_config.containers_config.agent_ip = "192.168.1.3"
+        s.attacker_obs_state.agent_reachable = ["192.168.1.1"]
+
+        result = ConnectionUtil._ssh_setup_connection(
+            a=a, credentials=credentials, proxy_connections=proxy_connections, s=s
+        )
+        mock_SSHClient.assert_not_called()
+        assert result
+
+    @patch("csle_common.util.emulation_util.EmulationUtil.execute_ssh_cmd")
+    def test_ssh_finalize_connection(self, mock_execute_ssh_cmd) -> None:
+        """
+        Test the helper function for finalizing a SSH connection and setting up the DTO
+
+        :param mock_execute_ssh_cmd: mock_execute_ssh_cmd
+
+        :return: None
+        """
+        mock_execute_ssh_cmd.side_effect = [
+            (b"output", b"error", 0.1),
+            (b"output", b"error", 0.1),
+            (b"output", b"error", 0.1),
+            (b"output", b"error", 0.1),
+            (b"(ALL) NOPASSWD: ALL", b"", 0.1),
+        ]
+
+        target_machine = MagicMock()
+        connection_setup_dto = MagicMock()
+        connection_setup_dto.target_connections = [MagicMock()]
+        connection_setup_dto.credentials = [MagicMock()]
+        connection_setup_dto.credentials[0].username = "user"
+        connection_setup_dto.ports = [22]
+        connection_setup_dto.proxies = [MagicMock()]
+        connection_setup_dto.ip = "192.168.1.2"
+
+        root, total_time = ConnectionUtil._ssh_finalize_connection(
+            target_machine=target_machine, connection_setup_dto=connection_setup_dto, i=0
+        )
+        assert not root
+        assert total_time
+        mock_execute_ssh_cmd.assert_called()
+
+    @patch("telnetlib.Telnet")
+    def test_telnet_setup_connection(self, mock_telnet) -> None:
+        """
+        Test the helper function for setting up a Telnet connection to a target machine
+
+        :param mock_telnet: mock_telnet
+
+        :return: None
+        """
+        mock_telnet_conn = MagicMock()
+        mock_telnet.return_value = mock_telnet_conn
+        mock_telnet_conn.read_until.side_effect = [
+            constants.TELNET.LOGIN_PROMPT,
+            constants.TELNET.PASSWORD_PROMPT,
+            constants.TELNET.PROMPT,
+        ]
+        mock_telnet_conn.write.return_value = None
+        mock_telnet_conn.read_until.return_value = b"$"
+
+        a = MagicMock()
+        a.ips = ["192.168.1.2"]
+        a.ips_match.side_effect = lambda ips: True
+
+        credentials = [MagicMock()]
+        credentials[0].service = constants.TELNET.SERVICE_NAME
+        credentials[0].username = "user"
+        credentials[0].pw = "password"
+        credentials[0].port = 23
+
+        proxy_connections = [MagicMock()]
+        proxy_connections[0].ip = "192.168.1.1"
+        proxy_connections[0].conn.get_transport.return_value = MagicMock()
+
+        s = MagicMock()
+        s.emulation_env_config.containers_config.agent_ip = "192.168.1.3"
+        s.attacker_obs_state.agent_reachable = ["192.168.1.1"]
+        s.get_attacker_machine.return_value.reachable = ["192.168.1.2"]
+        s.emulation_env_config.get_port_forward_port.return_value = 9999
+
+        result = ConnectionUtil._telnet_setup_connection(
+            a=a, credentials=credentials, proxy_connections=proxy_connections, s=s
+        )
+        mock_telnet.assert_called()
+        assert result
+
+    def test_telnet_finalize_connection(self) -> None:
+        """
+        Test the helper function for finalizing a Telnet connection to a target machine and creating the DTO
+
+        :return: None
+        """
+        target_machine = MagicMock()
+        connection_setup_dto = MagicMock()
+        mock_telnet_conn = [MagicMock() for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)]
+        connection_setup_dto.target_connections = mock_telnet_conn
+        connection_setup_dto.credentials = [
+            MagicMock() for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)
+        ]
+        for credential in connection_setup_dto.credentials:
+            credential.username = "user"
+        connection_setup_dto.tunnel_threads = [
+            MagicMock() for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)
+        ]
+        connection_setup_dto.forward_ports = [9999 for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)]
+        connection_setup_dto.ports = [23 for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)]
+        connection_setup_dto.proxies = [MagicMock() for _ in range(constants.ENV_CONSTANTS.ATTACKER_RETRY_CHECK_ROOT)]
+        connection_setup_dto.ip = "192.168.1.2"
+
+        for i, conn in enumerate(mock_telnet_conn):
+            conn.read_until.side_effect = [b"user may not run sudo" if i < 4 else b"(ALL) NOPASSWD: ALL"]
+
+        root, total_time = ConnectionUtil._telnet_finalize_connection(
+            target_machine=target_machine, i=0, connection_setup_dto=connection_setup_dto
+        )
+        assert not root
+        assert total_time
+
+    @patch("ftplib.FTP")
+    def test_ftp_setup_connection(self, mock_ftp) -> None:
+        """
+        Test the helper function for setting up a FTP connection
+
+        :param mock_ftp: mock_ftp
+
+        :return: None
+        """
+        mock_ftp_conn = MagicMock()
+        mock_ftp.return_value = mock_ftp_conn
+        mock_ftp_conn.login.return_value = "230 Login successful."
+
+        a = MagicMock()
+        a.ips = ["192.168.1.2"]
+        a.ips_match.side_effect = lambda ips: True
+
+        credentials = [MagicMock()]
+        credentials[0].service = constants.FTP.SERVICE_NAME
+        credentials[0].username = "user"
+        credentials[0].pw = "password"
+        credentials[0].port = 21
+
+        proxy_connections = [MagicMock()]
+        proxy_connections[0].ip = "192.168.1.1"
+        proxy_connections[0].conn.get_transport.return_value = MagicMock()
+        proxy_connections[0].conn.invoke_shell.return_value = MagicMock()
+
+        s = MagicMock()
+        s.emulation_env_config.containers_config.agent_ip = "192.168.1.3"
+        s.attacker_obs_state.agent_reachable = ["192.168.1.1"]
+        s.get_attacker_machine.return_value.reachable = ["192.168.1.2"]
+        s.emulation_env_config.get_port_forward_port.return_value = 9999
+
+        result = ConnectionUtil._ftp_setup_connection(
+            a=a, credentials=credentials, proxy_connections=proxy_connections, s=s
+        )
+        mock_ftp.assert_not_called()
+        assert result
+
+    def test_ftp_finalize_connection(self) -> None:
+        """
+        Test helper function for creating the connection DTO for FTP
+
+        :return: None
+        """
+        target_machine = MagicMock()
+        target_machine.ftp_connections = []
+
+        connection_setup_dto = MagicMock()
+        connection_setup_dto.target_connections = [MagicMock()]
+        connection_setup_dto.credentials = [MagicMock()]
+        connection_setup_dto.tunnel_threads = [MagicMock()]
+        connection_setup_dto.forward_ports = [9999]
+        connection_setup_dto.ports = [21]
+        connection_setup_dto.interactive_shells = [MagicMock()]
+        connection_setup_dto.proxies = [MagicMock()]
+        connection_setup_dto.ip = "192.168.1.2"
+
+        root, cost = ConnectionUtil._ftp_finalize_connection(
+            target_machine=target_machine, i=0, connection_setup_dto=connection_setup_dto
+        )
+        assert not root
+        assert cost == 0
+
+    def test__find_jump_host_connection(self) -> None:
+        """
+        Test utility function for finding a jump-host from the set of compromised machines to reach a target IP
+
+        :return: None
+        """
+        ip = "192.168.1.100"
+        s = MagicMock()
+        s.attacker_obs_state = MagicMock()
+        s.attacker_obs_state.agent_reachable = ["192.168.1.100"]
+        s.emulation_env_config.containers_config.agent_ip = "192.168.1.1"
+        s.emulation_env_config.get_hacker_connection.return_value = MagicMock()
+
+        result = ConnectionUtil.find_jump_host_connection(ip=ip, s=s)
+        assert result.ip == "192.168.1.1"
+        assert result.root
+
+    @patch("csle_common.util.emulation_util.EmulationUtil.execute_ssh_cmd")
+    def test_test_connection(self, mock_execute_ssh_cmd) -> None:
+        """
+        Test utility function for testing if a connection is alive or not
+
+        :param mock_execute_ssh_cmd: mock_execute_ssh_cmd
+
+        :return: None
+        """
+        mock_conn = MagicMock()
+        mock_execute_ssh_cmd.return_value = (b"user", b"", 0.1)
+        c = EmulationConnectionObservationState(
+            conn=mock_conn, credential=MagicMock(), root=False, service="ssh", port=22
+        )
+        result = ConnectionUtil.test_connection(c=c)
+        assert result
+
+    @patch("paramiko.SSHClient")
+    def test_reconnect_ssh(self, mock_SSHClient) -> None:
+        """
+        Test the method that reconnects the given SSH connection if it has died for some reason
+
+        :param mock_SSHClient: mock_SSHClient
+
+        :return: None
+        """
+        mock_credential = MagicMock()
+        mock_credential.username = "user"
+        mock_credential.pw = "password"
+        mock_credential.port = 22
+
+        c = EmulationConnectionObservationState(
+            conn=MagicMock(), credential=mock_credential, root=False, service="ssh", port=22, ip="192.168.1.2"
+        )
+        mock_ssh_client = mock_SSHClient.return_value
+        mock_transport = MagicMock()
+        mock_ssh_client.get_transport.return_value = mock_transport
+
+        result = ConnectionUtil.reconnect_ssh(c)
+        assert result.conn == mock_ssh_client
+
+    @patch("csle_common.util.connection_util.ConnectionUtil.reconnect_ssh")
+    @patch("telnetlib.Telnet")
+    def test_reconnect_telnet(self, mock_telnet, mock_reconnect_ssh) -> None:
+        """
+        Test the method that reconnects the given Telnet connection if it has died for some reason
+
+        :param mock_telnet: mock_telnet
+        :param mock_reconnect_ssh: mock_reconnect_ssh
+
+        :return: None
+        """
+        mock_credential = MagicMock()
+        mock_credential.username = "user"
+        mock_credential.pw = "password"
+
+        mock_proxy_conn = MagicMock()
+        proxy = EmulationConnectionObservationState(
+            conn=mock_proxy_conn, credential=mock_credential, root=False, service="ssh", port=22, ip="192.168.1.1"
+        )
+        c = EmulationConnectionObservationState(
+            conn=None, credential=mock_credential, root=False, service="telnet", port=23, ip="192.168.1.2", proxy=proxy
+        )
+
+        mock_telnet_conn = mock_telnet.return_value
+        mock_telnet_conn.read_until.side_effect = [
+            constants.TELNET.LOGIN_PROMPT,
+            constants.TELNET.PASSWORD_PROMPT,
+            constants.TELNET.PROMPT,
+        ]
+
+        mock_reconnect_ssh.return_value = proxy
+        result = ConnectionUtil.reconnect_telnet(c)
+        assert result.conn == mock_telnet_conn