Move the instance conn poller to Coriolis core

Most providers use a copy of those helpers, which poll until
the instance can be accessed using SSH or WINRM.

For this reason, we'll move them to the Coriolis core utils.

coriolis/constants.py

@@ -274,6 +274,10 @@ DEFAULT_OS_TYPE = OS_TYPE_LINUX
 VALID_OS_TYPES = [OS_TYPE_BSD, OS_TYPE_LINUX,
                   OS_TYPE_OS_X, OS_TYPE_SOLARIS, OS_TYPE_WINDOWS]
 
+PROTOCOL_SSH = "ssh"
+# WinRM is the Microsoft implementation of WSMAN.
+PROTOCOL_WINRM = "winrm"
+
 TMP_DIRS_KEY = "__tmp_dirs"
 
 COMPRESSION_FORMAT_GZIP = "gzip"

coriolis/tests/test_utils.py

@@ -1274,6 +1274,134 @@ class UtilsTestCase(test_base.CoriolisBaseTestCase):
         self.assertRaises(exception.UnrecognizedWorkerInitSystem,
                           stop_svc_undecorated, self.mock_ssh, 'svc_name')
 
+    def _get_mock_conn_info(self):
+        return {
+            "ip": "1.2.3.4",
+            "port": 2222,
+            "username": "someUser",
+            "password": "pwned",
+            "pkey": mock.Mock(),
+        }
+
+    @mock.patch.object(utils, "connect_ssh")
+    @mock.patch.object(utils, "exec_ssh_cmd")
+    @mock.patch("time.sleep")
+    def test_poll_instance_ssh(
+        self,
+        mock_sleep,
+        mock_exec_ssh,
+        mock_connect,
+    ):
+        mock_ssh_conn = mock.Mock()
+        mock_connect.side_effect = [
+            Exception,
+            mock_ssh_conn,
+            mock_ssh_conn,
+        ]
+        mock_exec_ssh.side_effect = [
+            Exception,
+            mock.sentinel.stdout,
+        ]
+        poll_interval = 5
+
+        connection_info = self._get_mock_conn_info()
+        utils.poll_instance_until_reachable(
+            connection_info=connection_info,
+            protocol=constants.PROTOCOL_SSH,
+            timeout=30,
+            poll_interval=poll_interval,
+        )
+
+        mock_connect.assert_has_calls(
+            [
+                mock.call(
+                    hostname=connection_info["ip"],
+                    port=connection_info["port"],
+                    username=connection_info["username"],
+                    password=connection_info["password"],
+                    pkey=connection_info["pkey"],
+                )
+            ] * 2)
+        mock_exec_ssh.assert_has_calls(
+            [mock.call(mock_ssh_conn, "exit 0")] * 2)
+        mock_ssh_conn.close.assert_has_calls([mock.call()] * 2)
+        mock_sleep.assert_has_calls([mock.call(poll_interval)] * 2)
+
+    @mock.patch.object(utils, "connect_ssh")
+    @mock.patch.object(utils, "exec_ssh_cmd")
+    @mock.patch("time.sleep")
+    @mock.patch("time.time")
+    def test_poll_instance_ssh_timeout(
+        self,
+        mock_time,
+        mock_sleep,
+        mock_exec_ssh,
+        mock_connect,
+    ):
+        poll_interval = 5
+        mock_time.side_effect = [x * 10 for x in range(20)]
+        mock_connect.side_effect = IOError
+        connection_info = self._get_mock_conn_info()
+        self.assertRaises(
+            exception.CoriolisException,
+            utils.poll_instance_until_reachable,
+            connection_info=connection_info,
+            protocol=constants.PROTOCOL_SSH,
+            timeout=30,
+            poll_interval=poll_interval,
+        )
+
+    @mock.patch("coriolis.wsman.WSManConnection", new_callable=mock.Mock)
+    @mock.patch("time.sleep")
+    def test_poll_instance_winrm(
+        self,
+        mock_sleep,
+        mock_wsman,
+    ):
+        mock_conn = mock.Mock()
+        mock_conn.exec_ps_command.side_effect = [
+            Exception, mock.sentinel.stdout]
+        mock_wsman.from_connection_info.return_value = mock_conn
+        poll_interval = 5
+
+        connection_info = self._get_mock_conn_info()
+        utils.poll_instance_until_reachable(
+            connection_info=connection_info,
+            protocol=constants.PROTOCOL_WINRM,
+            timeout=30,
+            poll_interval=poll_interval,
+        )
+
+        mock_wsman.from_connection_info.assert_has_calls(
+            [mock.call(connection_info)] * 2, any_order=True)
+        mock_conn.exec_ps_command.assert_has_calls(
+            [mock.call("whoami")] * 2)
+        mock_sleep.assert_called_once_with(poll_interval)
+
+    @mock.patch("coriolis.wsman.WSManConnection", new_callable=mock.Mock)
+    @mock.patch("time.sleep")
+    @mock.patch("time.time")
+    def test_poll_instance_winrm_timeout(
+        self,
+        mock_time,
+        mock_sleep,
+        mock_wsman,
+    ):
+        poll_interval = 5
+        mock_time.side_effect = [x * 10 for x in range(20)]
+        mock_conn = mock.Mock()
+        mock_conn.exec_ps_command.side_effect = IOError
+        mock_wsman.from_connection_info.return_value = mock_conn
+        connection_info = self._get_mock_conn_info()
+        self.assertRaises(
+            exception.CoriolisException,
+            utils.poll_instance_until_reachable,
+            connection_info=connection_info,
+            protocol=constants.PROTOCOL_WINRM,
+            timeout=30,
+            poll_interval=poll_interval,
+        )
+
 
 @ddt.ddt
 class Grub2ConfigEditorTestCase(test_base.CoriolisBaseTestCase):

coriolis/utils.py

@@ -1120,3 +1120,100 @@ def start_thread(target, args=(), kwargs=None, daemon=True):
     )
     thread.start()
     return thread
+
+
+def _poll_instance_until_reachable_ssh(
+    connection_info: dict,
+    timeout: int = 600,
+    poll_interval: int = 10,
+):
+    start = time.time()
+    while (time.time() - start) < timeout:
+        try:
+            ssh = connect_ssh(
+                hostname=connection_info["ip"],
+                port=connection_info["port"],
+                username=connection_info["username"],
+                password=connection_info["password"],
+                pkey=connection_info["pkey"],
+            )
+            try:
+                # "exit 0" should work across platforms.
+                # "whoami" would also work.
+                exec_ssh_cmd(ssh, "exit 0")
+            finally:
+                ssh.close()
+            LOG.debug("Instance reachable: %s", connection_info["ip"])
+            return
+        except Exception as err:
+            LOG.debug(
+                f"Could not connect to remote instance: {str(err)}. "
+                f"Retrying, time left: {timeout - (time.time() - start)}."
+            )
+        time.sleep(poll_interval)
+
+    raise exception.CoriolisException(
+        f"Operation timed out after waiting {timeout}s for the instance to be "
+        f"accessible via SSH."
+    )
+
+
+def _poll_instance_until_reachable_winrm(
+    connection_info: dict,
+    timeout: int = 600,
+    poll_interval: int = 10,
+):
+    start = time.time()
+    while (time.time() - start) < timeout:
+        try:
+            # Avoid a circular dependency.
+            from coriolis import wsman
+            conn = wsman.WSManConnection.from_connection_info(connection_info)
+            conn.exec_ps_command("whoami")
+            return
+        except Exception as ex:
+            LOG.debug(
+                f"Could not conect to Windows host: {str(ex)}. "
+                f"Retrying, time left: {timeout - (time.time() - start)}."
+            )
+        time.sleep(poll_interval)
+
+    raise exception.CoriolisException(
+        f"Operation timed out after waiting {timeout}s for Windows host to "
+        f"be accessible via WinRM."
+    )
+
+
+def poll_instance_until_reachable(
+    connection_info: dict,
+    protocol: str = constants.PROTOCOL_SSH,
+    timeout: int = 600,
+    poll_interval: int = 10,
+) -> paramiko.SSHClient:
+    """Poll until a given instance becomes reachable.
+
+    :param connection_info: a dict containing the following keys:
+        * ip
+        * port
+        * username
+        * password
+        * pkey - Paramiko keypair
+    :param protocol: connection protocol, "ssh" or "winrm"
+    :param timeout: the maximum amount of time to wait
+    :param poll_interval: the amount of time to wait between retries
+    """
+    # TODO(lpetrut): consider including the connection protocol in the
+    # connection info. We'd have to modify a few schemas used during os
+    # morphing. We currently pick the protocol based on the OS type but
+    # we may want to use SSH on Windows as well.
+    if protocol == constants.PROTOCOL_SSH:
+        helper = _poll_instance_until_reachable_ssh
+    elif protocol == constants.PROTOCOL_WINRM:
+        helper = _poll_instance_until_reachable_winrm
+    else:
+        raise exception.InvalidInput(
+            f"Unsupported instance connection protocol: {protocol}")
+    return helper(
+        connection_info=connection_info,
+        timeout=timeout,
+        poll_interval=poll_interval)