#!/usr/bin/env python # remote-run - Runs a command on another machine, for testing -----*- python -*- # # This source file is part of the Swift.org open source project # # Copyright (c) 2018 Apple Inc. and the Swift project authors # Licensed under Apache License v2.0 with Runtime Library Exception # # See https://swift.org/LICENSE.txt for license information # See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors # # ---------------------------------------------------------------------------- from __future__ import print_function import argparse import os import posixpath import subprocess import sys def quote(arg): return repr(arg) class CommandRunner(object): def __init__(self): self.verbose = False self.dry_run = False @staticmethod def _dirnames(files): return sorted(set(posixpath.dirname(f) for f in files)) def popen(self, command, **kwargs): if self.verbose: print(' '.join(command), file=sys.stderr) if self.dry_run: return None return subprocess.Popen(command, **kwargs) def send(self, local_to_remote_files): # Prepare the remote directory structure. # FIXME: This could be folded into the sftp connection below. dirs_to_make = self._dirnames(local_to_remote_files.viewvalues()) self.run_remote(['/bin/mkdir', '-p'] + dirs_to_make) # Send the local files. sftp_commands = ("-put {0} {1}".format(quote(local_file), quote(remote_file)) for local_file, remote_file in local_to_remote_files.viewitems()) self.run_sftp(sftp_commands) def fetch(self, local_to_remote_files): # Prepare the local directory structure. dirs_to_make = self._dirnames(local_to_remote_files.viewkeys()) mkdir_command = ['/bin/mkdir', '-p'] + dirs_to_make if self.verbose: print(' '.join(mkdir_command), file=sys.stderr) if not self.dry_run: subprocess.check_call(mkdir_command) # Fetch the remote files. sftp_commands = ("-get {0} {1}".format(quote(remote_file), quote(local_file)) for local_file, remote_file in local_to_remote_files.viewitems()) self.run_sftp(sftp_commands) def run_remote(self, command, remote_env={}): env_strings = ['{0}={1}'.format(k,v) for k,v in remote_env.viewitems()] remote_invocation = self.remote_invocation( ['/usr/bin/env'] + env_strings + command) remote_proc = self.popen(remote_invocation, stdin=subprocess.PIPE, stdout=None, stderr=None) if self.dry_run: return _, _ = remote_proc.communicate() if remote_proc.returncode: # FIXME: We may still want to fetch the output files to see what # went wrong. sys.exit(remote_proc.returncode) def run_sftp(self, commands): sftp_proc = self.popen(self.sftp_invocation(), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=None) concatenated_commands = '\n'.join(commands) if self.verbose: print(concatenated_commands, file=sys.stderr) if self.dry_run: return _, _ = sftp_proc.communicate(concatenated_commands) if sftp_proc.returncode: sys.exit(sftp_proc.returncode) class RemoteCommandRunner(CommandRunner): def __init__(self, host, identity_path, ssh_options, config_file): if ':' in host: (self.remote_host, self.port) = host.rsplit(':', 1) else: self.remote_host = host self.port = None self.identity_path = identity_path self.ssh_options = ssh_options self.config_file = config_file def common_options(self, port_flag): port_option = [port_flag, self.port] if self.port else [] config_option = ['-F', self.config_file] if self.config_file else [] identity_option = ( ['-i', self.identity_path] if self.identity_path else []) # Interleave '-o' with each custom option. # From https://stackoverflow.com/a/8168526, # with explanatory help from # https://spapas.github.io/2016/04/27/python-nested-list-comprehensions/ extra_options = [arg for option in self.ssh_options for arg in ["-o", option]] return port_option + identity_option + config_option + extra_options def remote_invocation(self, command): return (['/usr/bin/ssh', '-n'] + self.common_options(port_flag='-p') + [self.remote_host, '--'] + [quote(arg) for arg in command]) def sftp_invocation(self): return (['/usr/bin/sftp', '-b', '-', '-q', '-r'] + self.common_options(port_flag='-P') + [self.remote_host]) class LocalCommandRunner(CommandRunner): def __init__(self, sftp_server_path): self.sftp_server_path = sftp_server_path def remote_invocation(self, command): return command def sftp_invocation(self): return ['/usr/bin/sftp', '-b', '-', '-q', '-D', self.sftp_server_path] def find_transfers(args, source_prefix, dest_prefix): if source_prefix.endswith(posixpath.sep): source_prefix = source_prefix[:-len(posixpath.sep)] return dict((arg, dest_prefix + arg[len(source_prefix):]) for arg in args if arg.startswith(source_prefix)) def collect_remote_env(local_env=os.environ, prefix='REMOTE_RUN_CHILD_'): return dict((key[len(prefix):], value) for key, value in local_env.items() if key.startswith(prefix)) def main(): parser = argparse.ArgumentParser() parser.add_argument('-v', '--verbose', action='store_true', dest='verbose', help='print commands as they are run') parser.add_argument('-n', '--dry-run', action='store_true', dest='dry_run', help="print the commands that would have been run, but " "don't actually run them") parser.add_argument('--remote-dir', required=True, metavar='PATH', help='(required) a writable temporary path on the ' 'remote machine') parser.add_argument('--input-prefix', help='arguments matching this prefix will be uploaded') parser.add_argument('--output-prefix', help='arguments matching this prefix will be both ' 'uploaded and downloaded') parser.add_argument('--remote-input-prefix', default='input', help='input arguments use this prefix on the remote ' 'machine') parser.add_argument('--remote-output-prefix', default='output', help='output arguments use this prefix on the remote ' 'machine') parser.add_argument('-i', '--identity', dest='identity', metavar='FILE', help='an SSH identity file (private key) to use') parser.add_argument('-F', '--config-file', dest='config_file', metavar='FILE', help='an SSH configuration file') parser.add_argument('-o', '--ssh-option', action='append', default=[], dest='ssh_options', metavar='OPTION', help='extra SSH config options (man ssh_config)') parser.add_argument('--debug-as-local', metavar='/PATH/TO/SFTP-SERVER', help='run commands locally instead of over SSH, for ' 'debugging purposes. The "host" argument is ' 'omitted.') parser.add_argument('host', help='the host to connect to, in the form ' '[user@]host[:port]') parser.add_argument('command', nargs=argparse.REMAINDER, help='the command to run', metavar='command...') args = parser.parse_args() if args.debug_as_local: runner = LocalCommandRunner(args.debug_as_local) args.command.insert(0, args.host) del args.host else: runner = RemoteCommandRunner(args.host, args.identity, args.ssh_options, args.config_file) runner.dry_run = args.dry_run runner.verbose = args.verbose or args.dry_run assert not args.remote_dir == '/' upload_files = dict() download_files = dict() remote_test_specific_dir = None if args.input_prefix: assert not args.remote_input_prefix.startswith("..") remote_dir = posixpath.join(args.remote_dir, args.remote_input_prefix) input_files = find_transfers(args.command, args.input_prefix, remote_dir) assert not any(upload_files.has_key(f) for f in input_files) upload_files.update(input_files) if args.output_prefix: assert not args.remote_output_prefix.startswith("..") remote_dir = posixpath.join(args.remote_dir, args.remote_output_prefix) test_files = find_transfers(args.command, args.output_prefix, remote_dir) assert not any(upload_files.has_key(f) for f in test_files) upload_files.update(test_files) assert not any(download_files.has_key(f) for f in test_files) download_files.update(test_files) remote_test_specific_dir = remote_dir if remote_test_specific_dir: assert remote_test_specific_dir.startswith(args.remote_dir) runner.run_remote(['/bin/rm', '-rf', remote_test_specific_dir]) if upload_files: runner.send(upload_files) remote_env = collect_remote_env() translated_command = [upload_files.get(arg, download_files.get(arg, arg)) for arg in args.command] runner.run_remote(translated_command, remote_env) if download_files: runner.fetch(download_files) if __name__ == "__main__": main()