remote-run

#!/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__ importprint_function

importargparse
importos
importposixpath
importsubprocess
importsys

defquote(arg):
returnrepr(arg)

classCommandRunner(object):
def__init__(self):
self.verbose=False
self.dry_run=False

@staticmethod
def_dirnames(files):
returnsorted(set(posixpath.dirname(f) forfinfiles))

defpopen(self, command, **kwargs):
ifself.verbose:
print(' '.join(command), file=sys.stderr)
ifself.dry_run:
returnNone
returnsubprocess.Popen(command, **kwargs)

defsend(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))
forlocal_file, remote_file
inlocal_to_remote_files.viewitems())
self.run_sftp(sftp_commands)

deffetch(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
ifself.verbose:
print(' '.join(mkdir_command), file=sys.stderr)
ifnotself.dry_run:
subprocess.check_call(mkdir_command)

# Fetch the remote files.
sftp_commands= ("-get {0} {1}".format(quote(remote_file), 
quote(local_file))
forlocal_file, remote_file
inlocal_to_remote_files.viewitems())
self.run_sftp(sftp_commands)

defrun_remote(self, command, remote_env={}):
env_strings= ['{0}={1}'.format(k,v) fork,vinremote_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)
ifself.dry_run:
return
_, _=remote_proc.communicate()
ifremote_proc.returncode:
# FIXME: We may still want to fetch the output files to see what
# went wrong.
sys.exit(remote_proc.returncode)

defrun_sftp(self, commands):
sftp_proc=self.popen(self.sftp_invocation(), stdin=subprocess.PIPE,
stdout=subprocess.PIPE, stderr=None)
concatenated_commands='\n'.join(commands)
ifself.verbose:
print(concatenated_commands, file=sys.stderr)
ifself.dry_run:
return
_, _=sftp_proc.communicate(concatenated_commands)
ifsftp_proc.returncode:
sys.exit(sftp_proc.returncode)

classRemoteCommandRunner(CommandRunner):
def__init__(self, host, identity_path, ssh_options, config_file):
if':'inhost:
 (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

defcommon_options(self, port_flag):
port_option= [port_flag, self.port] ifself.portelse []
config_option= ['-F', self.config_file] ifself.config_fileelse []
identity_option= (
 ['-i', self.identity_path] ifself.identity_pathelse [])
# 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= [argforoptioninself.ssh_options
forargin ["-o", option]]
returnport_option+identity_option+config_option+extra_options

defremote_invocation(self, command):
return (['/usr/bin/ssh', '-n'] +
self.common_options(port_flag='-p') +
 [self.remote_host, '--'] +
 [quote(arg) forargincommand])

defsftp_invocation(self):
return (['/usr/bin/sftp', '-b', '-', '-q', '-r'] +
self.common_options(port_flag='-P') +
 [self.remote_host])

classLocalCommandRunner(CommandRunner):
def__init__(self, sftp_server_path):
self.sftp_server_path=sftp_server_path

defremote_invocation(self, command):
returncommand

defsftp_invocation(self):
return ['/usr/bin/sftp', '-b', '-', '-q', '-D', self.sftp_server_path]

deffind_transfers(args, source_prefix, dest_prefix):
ifsource_prefix.endswith(posixpath.sep):
source_prefix=source_prefix[:-len(posixpath.sep)]
returndict((arg, dest_prefix+arg[len(source_prefix):])
forarginargsifarg.startswith(source_prefix))

defcollect_remote_env(local_env=os.environ, prefix='REMOTE_RUN_CHILD_'):
returndict((key[len(prefix):], value)
forkey, valueinlocal_env.items() ifkey.startswith(prefix))

defmain():
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()

ifargs.debug_as_local:
runner=LocalCommandRunner(args.debug_as_local)
args.command.insert(0, args.host)
delargs.host
else:
runner=RemoteCommandRunner(args.host,
args.identity,
args.ssh_options,
args.config_file)
runner.dry_run=args.dry_run
runner.verbose=args.verboseorargs.dry_run

assertnotargs.remote_dir=='/'

upload_files=dict()
download_files=dict()
remote_test_specific_dir=None
ifargs.input_prefix:
assertnotargs.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)
assertnotany(upload_files.has_key(f) forfininput_files)
upload_files.update(input_files)
ifargs.output_prefix:
assertnotargs.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)
assertnotany(upload_files.has_key(f) forfintest_files)
upload_files.update(test_files)
assertnotany(download_files.has_key(f) forfintest_files)
download_files.update(test_files)
remote_test_specific_dir=remote_dir

ifremote_test_specific_dir:
assertremote_test_specific_dir.startswith(args.remote_dir)
runner.run_remote(['/bin/rm', '-rf', remote_test_specific_dir])
ifupload_files:
runner.send(upload_files)

remote_env=collect_remote_env()

translated_command= [upload_files.get(arg, download_files.get(arg, arg))
forarginargs.command]
runner.run_remote(translated_command, remote_env)

ifdownload_files:
runner.fetch(download_files)

if__name__=="__main__":
main()