diff options
author | Jonathan Abrahams <jonathan@mongodb.com> | 2017-08-04 11:50:11 -0400 |
---|---|---|
committer | Jonathan Abrahams <jonathan@mongodb.com> | 2017-08-04 11:50:11 -0400 |
commit | b90b6c5eb16f3bdc974e26c4ef0f73488750d17d (patch) | |
tree | 9f43a42dcf687cc6077e08157b4c0f3ee3cda2b7 /buildscripts/remote_operations.py | |
parent | 4f07e8245a8b41a9b6a143cc39bb02c968004252 (diff) | |
download | mongo-b90b6c5eb16f3bdc974e26c4ef0f73488750d17d.tar.gz |
SERVER-30146 Evergreen scripts for remote hosts: ssh & scp
Diffstat (limited to 'buildscripts/remote_operations.py')
-rwxr-xr-x | buildscripts/remote_operations.py | 324 |
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() |