SERVER-30146 Evergreen scripts for remote hosts: ssh & scp

1 files changed, 324 insertions, 0 deletions

diff --git a/buildscripts/remote_operations.py b/buildscripts/remote_operations.py
new file mode 100755
index 00000000000..162f45629dc
--- /dev/null
+++ b/buildscripts/remote_operations.py
@@ -0,0 +1,324 @@
+#!/usr/bin/env python
+
+"""Remote access utilities, via ssh & scp."""
+
+from __future__ import print_function
+
+import optparse
+import os
+import re
+import shlex
+import sys
+import time
+
+# The subprocess32 module is untested on Windows and thus isn't recommended for use, even when it's
+# installed. See https://github.com/google/python-subprocess32/blob/3.2.7/README.md#usage.
+if os.name == "posix" and sys.version_info[0] == 2:
+    try:
+        import subprocess32 as subprocess
+    except ImportError:
+        import warnings
+        warnings.warn(("Falling back to using the subprocess module because subprocess32 isn't"
+                       " available. When using the subprocess module, a child process may trigger"
+                       " an invalid free(). See SERVER-22219 for more details."),
+                      RuntimeWarning)
+        import subprocess
+else:
+    import subprocess
+
+# Get relative imports to work when the package is not installed on the PYTHONPATH.
+if __name__ == "__main__" and __package__ is None:
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+_OPERATIONS = ["shell", "copy_to", "copy_from"]
+
+
+class RemoteOperations(object):
+    """Class to support remote operations."""
+
+    def __init__(self,
+                 user_host,
+                 ssh_options="",
+                 retries=0,
+                 retry_sleep=0,
+                 debug=False,
+                 use_shell=False):
+
+        self.user_host = user_host
+        self.ssh_options = ssh_options if ssh_options else ""
+        self.retries = retries
+        self.retry_sleep = retry_sleep
+        self.debug = debug
+        self.use_shell = use_shell
+        # Check if we can remotely access the host.
+        self._access_code, self._access_buff = self._remote_access()
+
+    def _call(self, cmd):
+        if self.debug:
+            print(cmd)
+        # If use_shell is False we need to split the command up into a list.
+        if not self.use_shell:
+            cmd = shlex.split(cmd)
+        # Use a common pipe for stdout & stderr for logging.
+        process = subprocess.Popen(cmd,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.STDOUT,
+                                   shell=self.use_shell)
+        buff_stdout, buff_stderr = process.communicate()
+        return process.poll(), buff_stdout
+
+    def _remote_access(self):
+        """ This will check if a remote session is possible. """
+        cmd = "ssh {} {} date".format(self.ssh_options, self.user_host)
+        attempt_num = 0
+        buff = ""
+        while True:
+            ret, buff = self._call(cmd)
+            if not ret:
+                return ret, buff
+            attempt_num += 1
+            if attempt_num > self.retries:
+                break
+            if self.debug:
+                print("Failed remote attempt {}, retrying in {} seconds".format(
+                    attempt_num, self.retry_sleep))
+            time.sleep(self.retry_sleep)
+        return ret, buff
+
+    def _perform_operation(self, cmd):
+        return self._call(cmd)
+
+    def access_established(self):
+        """ Returns True if initial access was establsished. """
+        return not self._access_code
+
+    def access_info(self):
+        """ Returns return code and output buffer from initial access attempt(s). """
+        return self._access_code, self._access_buff
+
+    def operation(self, operation_type, operation_param, operation_dir=None):
+        """ Main entry for remote operations. Returns (code, output).
+
+            'operation_type' supports remote shell and copy operations.
+            'operation_param' can either be a list or string of commands or files.
+            'operation_dir' is '.' if unspecified for 'copy_*'.
+        """
+
+        if not self.access_established():
+            return self.access_info()
+
+        # File names with a space must be quoted, since we permit the
+        # the file names to be either a string or a list.
+        if operation_type != "shell" and isinstance(operation_param, str):
+            operation_param = shlex.split(operation_param)
+
+        cmds = []
+        if operation_type == "shell":
+            if operation_dir is not None:
+                operation_param = "cd {}; {}".format(operation_dir, operation_param)
+            cmd = "ssh {} {} '{}'".format(self.ssh_options, self.user_host, operation_param)
+            cmds.append(cmd)
+
+        elif operation_type == "copy_to":
+            cmd = "scp -r {}".format(self.ssh_options)
+            # To support spaces in the filename or directory, we quote them one at a time.
+            for file in operation_param:
+                cmd += "\"{}\" ".format(file)
+            operation_dir = operation_dir if operation_dir else ""
+            cmd += " {}:{}".format(self.user_host, operation_dir)
+            cmds.append(cmd)
+
+        elif operation_type == "copy_from":
+            operation_dir = operation_dir if operation_dir else "."
+            if not os.path.isdir(operation_dir):
+                raise ValueError(
+                    "Local directory '{}' does not exist.".format(operation_dir))
+
+            # We support multiple files being copied from the remote host
+            # by invoking scp for each file specified.
+            # Note - this is a method which scp does not support directly.
+            for file in operation_param:
+                cmd = "scp -r {} {}:".format(self.ssh_options, self.user_host)
+                # Quote, and escape the file if there are spaces.
+                # Note - we do not support other non-ASCII characters in a file name.
+                if " " in file:
+                    file = re.escape("\"{}\"".format(file))
+                cmd += "{} {}".format(file, operation_dir)
+                cmds.append(cmd)
+
+        else:
+            raise ValueError(
+                "Invalid operation '{}' specified, choose from {}.".format(
+                    operation_type, _OPERATIONS))
+
+        final_ret = 0
+        for cmd in cmds:
+            ret, buff = self._perform_operation(cmd)
+            buff += buff
+            final_ret = final_ret or ret
+
+        return final_ret, buff
+
+    def shell(self, operation_param, operation_dir=None):
+        """ Helper for remote shell operations. """
+        return self.operation(
+            operation_type="shell",
+            operation_param=operation_param,
+            operation_dir=operation_dir)
+
+    def copy_to(self, operation_param, operation_dir=None):
+        """ Helper for remote copy_to operations. """
+        return self.operation(
+            operation_type="copy_to",
+            operation_param=operation_param,
+            operation_dir=operation_dir)
+
+    def copy_from(self, operation_param, operation_dir=None):
+        """ Helper for remote copy_from operations. """
+        return self.operation(
+            operation_type="copy_from",
+            operation_param=operation_param,
+            operation_dir=operation_dir)
+
+
+def main():
+
+    parser = optparse.OptionParser(description=__doc__)
+    control_options = optparse.OptionGroup(parser, "Control options")
+    shell_options = optparse.OptionGroup(parser, "Shell options")
+    copy_options = optparse.OptionGroup(parser, "Copy options")
+
+    parser.add_option("--userHost",
+                      dest="user_host",
+                      default=None,
+                      help="User and remote host to execute commands on [REQUIRED]."
+                           " Examples, 'user@1.2.3.4' or 'user@myhost.com'.")
+
+    parser.add_option("--operation",
+                      dest="operation",
+                      default="shell",
+                      choices=_OPERATIONS,
+                      help="Remote operation to perform, choose one of '{}',"
+                           " defaults to '%default'.".format(", ".join(_OPERATIONS)))
+
+    control_options.add_option("--sshOptions",
+                               dest="ssh_options",
+                               default=None,
+                               action="append",
+                               help="SSH connection options."
+                                    " More than one option can be specified either"
+                                    " in one quoted string or by specifying"
+                                    " this option more than once. Example options:"
+                                    " '-i $HOME/.ssh/access.pem -o ConnectTimeout=10"
+                                    " -o ConnectionAttempts=10'")
+
+    control_options.add_option("--retries",
+                               dest="retries",
+                               type=int,
+                               default=0,
+                               help="Number of retries to attempt for operation,"
+                                    " defaults to '%default'.")
+
+    control_options.add_option("--retrySleep",
+                               dest="retry_sleep",
+                               type=int,
+                               default=10,
+                               help="Number of seconds to wait between retries,"
+                                    " defaults to '%default'.")
+
+    control_options.add_option("--debug",
+                               dest="debug",
+                               action="store_true",
+                               default=False,
+                               help="Provides debug output.")
+
+    control_options.add_option("--verbose",
+                               dest="verbose",
+                               action="store_true",
+                               default=False,
+                               help="Print exit status and output at end.")
+
+    shell_options.add_option("--commands",
+                             dest="remote_commands",
+                             default=None,
+                             action="append",
+                             help="Commands to excute on the remote host. The"
+                                  " commands must be separated by a ';' and can either"
+                                  " be specifed in a quoted string or by specifying"
+                                  " this option more than once. A ';' will be added"
+                                  " between commands when this option is specifed"
+                                  " more than once.")
+
+    shell_options.add_option("--commandDir",
+                             dest="command_dir",
+                             default=None,
+                             help="Working directory on remote to execute commands"
+                                  " form. Defaults to remote login directory.")
+
+    copy_options.add_option("--file",
+                            dest="files",
+                            default=None,
+                            action="append",
+                            help="The file to copy to/from remote host. To"
+                                 " support spaces in the file, each file must be"
+                                 " specified using this option more than once.")
+
+    copy_options.add_option("--remoteDir",
+                            dest="remote_dir",
+                            default=None,
+                            help="Remote directory to copy to, only applies when"
+                                 " operation is 'copy_to'. Defaults to the login"
+                                 " directory on the remote host.")
+
+    copy_options.add_option("--localDir",
+                            dest="local_dir",
+                            default=".",
+                            help="Local directory to copy to, only applies when"
+                                 " operation is 'copy_from'. Defaults to the"
+                                 " current directory, '%default'.")
+
+    parser.add_option_group(control_options)
+    parser.add_option_group(shell_options)
+    parser.add_option_group(copy_options)
+
+    (options, args) = parser.parse_args()
+
+    if not getattr(options, "user_host", None):
+        parser.print_help()
+        parser.error("Missing required option")
+
+    if options.operation == "shell":
+        if not getattr(options, "remote_commands", None):
+            parser.print_help()
+            parser.error("Missing required '{}' option '{}'".format(
+                options.operation, "--commands"))
+        operation_param = ";".join(options.remote_commands)
+        operation_dir = options.command_dir
+    else:
+        if not getattr(options, "files", None):
+            parser.print_help()
+            parser.error("Missing required '{}' option '{}'".format(
+                options.operation, "--file"))
+        operation_param = options.files
+        if options.operation == "copy_to":
+            operation_dir = options.remote_dir
+        else:
+            operation_dir = options.local_dir
+
+    ssh_options = None if not options.ssh_options else " ".join(options.ssh_options)
+    remote_op = RemoteOperations(
+        user_host=options.user_host,
+        ssh_options=ssh_options,
+        retries=options.retries,
+        retry_sleep=options.retry_sleep,
+        debug=options.debug)
+    ret_code, buffer = remote_op.operation(options.operation, operation_param, operation_dir)
+    if options.verbose:
+        print("Return code: {} for command {}".format(ret_code, sys.argv))
+        print(buffer)
+
+    sys.exit(ret_code)
+
+
+if __name__ == "__main__":
+    main()