diff options
| -rwxr-xr-x | buildscripts/remote_operations.py | 324 | ||||
| -rwxr-xr-x | buildscripts/tests/test_remote_operations.py | 344 |
2 files changed, 668 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() diff --git a/buildscripts/tests/test_remote_operations.py b/buildscripts/tests/test_remote_operations.py new file mode 100755 index 00000000000..cae02865d71 --- /dev/null +++ b/buildscripts/tests/test_remote_operations.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python + +"""Unit test for buildscripts/remote_operations.py. + + Note - Tests require sshd to be enabled on localhost with paswordless login.""" + +import os +import shutil +import sys +import tempfile +import unittest + +if __name__ == "__main__" and __package__ is None: + sys.path.append(os.getcwd()) +from buildscripts import remote_operations as rop + + +class RemoteOperationsTestCase(unittest.TestCase): + def setUp(self): + self.temp_local_dir = tempfile.mkdtemp() + self.temp_remote_dir = tempfile.mkdtemp() + self.rop = rop.RemoteOperations(user_host="localhost") + self.rop_shell = rop.RemoteOperations(user_host="localhost", use_shell=True) + + def tearDown(self): + shutil.rmtree(self.temp_local_dir) + shutil.rmtree(self.temp_remote_dir) + + +class RemoteOperationConnection(RemoteOperationsTestCase): + def runTest(self): + + self.assertTrue(self.rop.access_established()) + ret, buff = self.rop.access_info() + self.assertEqual(0, ret) + + # Invalid host + remote_op = rop.RemoteOperations(user_host="badhost") + ret, buff = remote_op.access_info() + self.assertFalse(remote_op.access_established()) + self.assertEqual(255, ret) + self.assertIsNotNone(buff) + + # Invalid host with retries + remote_op = rop.RemoteOperations(user_host="badhost2", retries=3) + ret, buff = remote_op.access_info() + self.assertFalse(remote_op.access_established()) + self.assertNotEqual(0, ret) + self.assertIsNotNone(buff) + + # Invalid host with retries & retry_sleep + remote_op = rop.RemoteOperations(user_host="badhost3", retries=3, retry_sleep=1) + ret, buff = remote_op.access_info() + self.assertFalse(remote_op.access_established()) + self.assertNotEqual(0, ret) + self.assertIsNotNone(buff) + + # Valid host with invalid ssh_options + ssh_options = "-o invalid" + remote_op = rop.RemoteOperations(user_host="localhost", ssh_options=ssh_options) + ret, buff = remote_op.access_info() + self.assertFalse(remote_op.access_established()) + self.assertNotEqual(0, ret) + self.assertIsNotNone(buff) + + # Valid host with valid ssh_options + ssh_options = "-v -o ConnectTimeout=10 -o ConnectionAttempts=10" + remote_op = rop.RemoteOperations(user_host="localhost", ssh_options=ssh_options) + ret, buff = remote_op.access_info() + self.assertTrue(remote_op.access_established()) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + +class RemoteOperationShell(RemoteOperationsTestCase): + def runTest(self): + + # Shell connect + ret, buff = self.rop.shell("uname") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop_shell.shell("uname") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop.operation("shell", "uname") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + # Invalid command + ret, buff = self.rop.shell("invalid_command") + self.assertNotEqual(0, ret) + self.assertIsNotNone(buff) + + # Multiple commands + ret, buff = self.rop.shell("date; whoami; ls") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop_shell.shell("date; whoami; ls") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + # Command with single quotes + ret, buff = self.rop.shell("echo 'hello there' | grep 'hello'") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop_shell.shell("echo 'hello there'| grep 'hello'") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + # Command with escaped double quotes + ret, buff = self.rop.shell("echo \"hello there\" | grep \"hello\"") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop_shell.shell("echo \"hello there\" | grep \"hello\"") + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + # Command with directory and pipe + ret, buff = self.rop.shell( + "touch {dir}/{file}; ls {dir} | grep {file}".format( + file="b", + dir=self.temp_local_dir)) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + ret, buff = self.rop_shell.shell( + "touch {dir}/{file}; ls {dir} | grep {file}".format( + file="c", + dir=self.temp_local_dir)) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + + +class RemoteOperationCopy(RemoteOperationsTestCase): + def runTest(self): + + # Copy to remote + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop.copy_to(l_temp_path, self.temp_remote_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(self.temp_remote_dir, l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop_shell.copy_to(l_temp_path, self.temp_remote_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(self.temp_remote_dir, l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop.operation("copy_to", l_temp_path, self.temp_remote_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + self.assertTrue(os.path.isfile(r_temp_path)) + + # Copy multiple files to remote + num_files = 3 + l_temp_files = [] + for i in range(num_files): + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + l_temp_files.append(l_temp_path) + ret, buff = self.rop.copy_to(" ".join(l_temp_files), self.temp_remote_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + r_temp_path = os.path.join(self.temp_remote_dir, os.path.basename(l_temp_files[i])) + self.assertTrue(os.path.isfile(r_temp_path)) + + num_files = 3 + l_temp_files = [] + for i in range(num_files): + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + l_temp_files.append(l_temp_path) + ret, buff = self.rop_shell.copy_to(" ".join(l_temp_files), self.temp_remote_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + r_temp_path = os.path.join(self.temp_remote_dir, os.path.basename(l_temp_files[i])) + self.assertTrue(os.path.isfile(r_temp_path)) + + # Copy to remote without directory + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop.copy_to(l_temp_path) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(os.environ["HOME"], l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + os.remove(r_temp_path) + + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir)[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop_shell.copy_to(l_temp_path) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(os.environ["HOME"], l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + os.remove(r_temp_path) + + # Copy to remote with space in file name, note it must be quoted. + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir, prefix="filename with space")[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop.copy_to("'{}'".format(l_temp_path)) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(os.environ["HOME"], l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + os.remove(r_temp_path) + + l_temp_path = tempfile.mkstemp(dir=self.temp_local_dir, prefix="filename with space")[1] + l_temp_file = os.path.basename(l_temp_path) + ret, buff = self.rop_shell.copy_to("'{}'".format(l_temp_path)) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + r_temp_path = os.path.join(os.environ["HOME"], l_temp_file) + self.assertTrue(os.path.isfile(r_temp_path)) + os.remove(r_temp_path) + + # Copy from remote + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + ret, buff = self.rop.copy_from(r_temp_path, self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + l_temp_path = os.path.join(self.temp_local_dir, r_temp_file) + self.assertTrue(os.path.isfile(l_temp_path)) + + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + ret, buff = self.rop_shell.copy_from(r_temp_path, self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + l_temp_path = os.path.join(self.temp_local_dir, r_temp_file) + self.assertTrue(os.path.isfile(l_temp_path)) + + # Copy from remote without directory + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + ret, buff = self.rop.copy_from(r_temp_path) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + self.assertTrue(os.path.isfile(r_temp_file)) + os.remove(r_temp_file) + + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + ret, buff = self.rop_shell.copy_from(r_temp_path) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + self.assertTrue(os.path.isfile(r_temp_file)) + os.remove(r_temp_file) + + # Copy from remote with space in file name, note it must be quoted. + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir, prefix="filename with space")[1] + r_temp_file = os.path.basename(r_temp_path) + ret, buff = self.rop.copy_from("'{}'".format(r_temp_path)) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + self.assertTrue(os.path.isfile(r_temp_file)) + os.remove(r_temp_file) + + # Copy multiple files from remote + num_files = 3 + r_temp_files = [] + for i in range(num_files): + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + r_temp_files.append(r_temp_path) + ret, buff = self.rop.copy_from(" ".join(r_temp_files), self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + basefile_name = os.path.basename(l_temp_files[i]) + l_temp_path = os.path.join(self.temp_local_dir, basefile_name) + self.assertTrue(os.path.isfile(l_temp_path)) + + num_files = 3 + r_temp_files = [] + for i in range(num_files): + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir)[1] + r_temp_file = os.path.basename(r_temp_path) + r_temp_files.append(r_temp_path) + ret, buff = self.rop_shell.copy_from(" ".join(r_temp_files), self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + basefile_name = os.path.basename(l_temp_files[i]) + l_temp_path = os.path.join(self.temp_local_dir, basefile_name) + self.assertTrue(os.path.isfile(l_temp_path)) + + # Copy files from remote with wilcard + num_files = 3 + r_temp_files = [] + for i in range(num_files): + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir, prefix="wild1")[1] + r_temp_file = os.path.basename(r_temp_path) + r_temp_files.append(r_temp_path) + r_temp_path = os.path.join(self.temp_remote_dir, "wild1*") + ret, buff = self.rop.copy_from(r_temp_path, self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + l_temp_path = os.path.join(self.temp_local_dir, os.path.basename(r_temp_files[i])) + self.assertTrue(os.path.isfile(l_temp_path)) + + num_files = 3 + r_temp_files = [] + for i in range(num_files): + r_temp_path = tempfile.mkstemp(dir=self.temp_remote_dir, prefix="wild2")[1] + r_temp_file = os.path.basename(r_temp_path) + r_temp_files.append(r_temp_path) + r_temp_path = os.path.join(self.temp_remote_dir, "wild2*") + ret, buff = self.rop_shell.copy_from(r_temp_path, self.temp_local_dir) + self.assertEqual(0, ret) + self.assertIsNotNone(buff) + for i in range(num_files): + l_temp_path = os.path.join(self.temp_local_dir, os.path.basename(r_temp_files[i])) + self.assertTrue(os.path.isfile(l_temp_path)) + + # Local directory does not exist. + self.assertRaises(ValueError, lambda: self.rop_shell.copy_from(r_temp_path, "bad_dir")) + + +class RemoteOperation(RemoteOperationsTestCase): + def runTest(self): + + # Invalid operation + self.assertRaises(ValueError, lambda: self.rop.operation("invalid", None)) + + +if __name__ == "__main__": + unittest.main() |