remote-run

#!/usr/bin/env python3
# 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
#
# ----------------------------------------------------------------------------

importargparse
importos
importposixpath
importsubprocess
importsys
importshutil

defquote(arg):
returnrepr(arg)

classCommandRunner(object):
def__init__(self):
self.verbose=False
self.dry_run=False
self.ignore_rsync_failure=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)

defmkdirs_remote(self, directories):
ifdirectories:
mkdir_command= ['/bin/mkdir', '-p'] +directories
self.run_remote(mkdir_command)

defsend(self, input_prefix, remote_prefix, local_to_remote_files):
# Prepare the remote directory structure.
self.mkdirs_remote([remote_prefix])

self.run_rsync_to(input_prefix, local_to_remote_files, remote_prefix)

deffetch(self, output_prefix, remote_prefix, remote_to_local_files):
# Prepare the local directory structure.
mkdir_command= ['/bin/mkdir', '-p', output_prefix]
ifself.verbose:
print(' '.join(mkdir_command), file=sys.stderr)
ifnotself.dry_run:
subprocess.check_call(mkdir_command)

self.run_rsync_from(remote_prefix, remote_to_local_files, output_prefix)

# Recover from random and transient errors that occur when SSHing to devices, in particular devicecompute
@staticmethod
defshould_remote_command_retry(stderr):
if"banner line contains invalid characters"instderr:
returnTrue

if"Connection to localhost closed by remote host"instderr:
returnTrue

if"kex_exchange_identification: Connection closed by remote host"instderr:
returnTrue

if"rsync error: unexplained error"instderr:
returnTrue

# Fallthrough. The error is not known and shouldn't be retried
returnFalse

defrun_remote(self, command, remote_env={}):
env_strings= ['{0}={1}'.format(k,v) fork,vinsorted(remote_env.items())]
remote_invocation=self.remote_invocation(
 ['/usr/bin/env'] +env_strings+command)
attempt=1
remote_proc=None

whileattempt<=3:
remote_proc=self.popen(
remote_invocation,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifself.dry_run:
return

stdout, stderr=remote_proc.communicate()
stdout=stdout.decode(encoding='utf-8', errors='replace')
stderr=stderr.decode(encoding='utf-8', errors='replace')

# Print stdout to screen
print(stdout, end='')

# This is a transient and random error, and if this occurs, we should simply retry our ssh command.
ifself.should_remote_command_retry(stderr):
attempt+=1
continue

print(stderr, end='', file=sys.stderr)

# Process error code
ifremote_proc.returncode:
# FIXME: We may still want to fetch the output files to see what
# went wrong.
sys.exit(remote_proc.returncode)
else:
# Nothing went wrong. Return
return

ifattempt>3andremote_procisnotNone:
sys.exit(remote_proc.returncode)

defrun_rsync(self, invocation, sources):
attempt=1
rsync_proc=None

whileattempt<=3:
rsync_proc=self.popen(
invocation,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifself.dry_run:
return

sources='\n'.join(sources)
ifself.verbose:
print(sources, file=sys.stderr)
stdout, stderr=rsync_proc.communicate(sources.encode('utf-8'))
stdout=stdout.decode(encoding='utf-8', errors='replace')
stderr=stderr.decode(encoding='utf-8', errors='replace')

# Print stdout to screen
print(stdout, end='')

# This is a transient and random error, and if this occurs, we should simply retry our ssh command.
ifself.should_remote_command_retry(stderr):
attempt+=1
continue

print(stderr, end='', file=sys.stderr)

# Process error code
ifrsync_proc.returncode:
ifself.ignore_rsync_failure:
return
else:
sys.exit(rsync_proc.returncode)
else:
# Nothing went wrong. Return
return

ifattempt>3andrsync_procisnotNone:
sys.exit(rsync_proc.returncode)

defrun_rsync_to(self, prefix, sources, dest):
self.run_rsync(self.rsync_to_invocation(prefix, dest), sources)

defrun_rsync_from(self, prefix, sources, dest):
self.run_rsync(self.rsync_from_invocation(prefix, dest), sources)

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])

defrsync_invocation(self, source, dest):
return ['/usr/bin/rsync', '-arRz', '--files-from=-',
'-e', ' '.join([quote(x) forxin
 ['/usr/bin/ssh'] +
self.common_options(port_flag='-p')]),
source,
dest]

defrsync_from_invocation(self, source, dest):
returnself.rsync_invocation(self.remote_host+':'+source, dest)

defrsync_to_invocation(self, source, dest):
returnself.rsync_invocation(source, self.remote_host+':'+dest)

classLocalCommandRunner(CommandRunner):
defremote_invocation(self, command):
returncommand

defrsync_invocation(self, source, dest):
return ['/usr/bin/rsync', '-arz', '--files-from=-', source, dest]

defrsync_from_invocation(self, source, dest):
returnself.rsync_invocation(source, dest)

defrsync_to_invocation(self, source, dest):
returnself.rsync_invocation(source, dest)

defstrip_sep(name):
lsep=len(posixpath.sep)
whilename.startswith(posixpath.sep):
name=name[lsep:]
returnname

classRemotePathSet(object):
def__init__(self):
self.inputs=set()
self.nodir_inputs=set()
self.existing_outputs=set()
self.outputs=set()
self.existing_nodir_outputs=set()
self.nodir_outputs=set()

classPrefixProcessor(object):
def__init__(self,
input_prefix, output_prefix,
remote_dir, remote_input_prefix,
remote_output_prefix,
path_set):
assertnotremote_input_prefix.startswith('..')
assertnotremote_output_prefix.startswith('..')

self.input_prefix=input_prefix
self.output_prefix=output_prefix
self.remote_dir=remote_dir
self.remote_input_prefix=remote_input_prefix
self.remote_output_prefix=remote_output_prefix
self.path_set=path_set

ifself.input_prefix:
whileself.input_prefix.endswith(posixpath.sep):
self.input_prefix=self.input_prefix[:-len(posixpath.sep)]
ifself.output_prefix:
whileself.output_prefix.endswith(posixpath.sep):
self.output_prefix=self.output_prefix[:-len(posixpath.sep)]

ifself.input_prefix:
split=posixpath.split(self.input_prefix)
self.input_prefix_split= (len(self.input_prefix), split[0], split[1])
else:
self.input_prefix_split= (0, '', '')

ifself.output_prefix:
split=posixpath.split(self.output_prefix)
self.output_prefix_split= (len(self.output_prefix),
split[0],
split[1])
else:
self.output_prefix_split= (0, '', '')

defprocess_one(self, orig_name):
iplen, ipfxdir, ipfx=self.input_prefix_split
oplen, opfxdir, opfx=self.output_prefix_split

ifiplenandorig_name.startswith(self.input_prefix):
name=orig_name[iplen:]
ifnotname.startswith(posixpath.sep):
name=ipfx+name
self.path_set.nodir_inputs.add(name)
else:
name=strip_sep(name)
self.path_set.inputs.add(name)
returnposixpath.join(self.remote_dir,
self.remote_input_prefix,
name)

ifoplenandorig_name.startswith(self.output_prefix):
name=orig_name[oplen:]
ifnotname.startswith(posixpath.sep):
name=opfx+name
self.path_set.nodir_outputs.add(name)
ifos.path.exists(orig_name):
self.path_set.existing_nodir_outputs.add(name)
else:
name=strip_sep(name)
self.path_set.outputs.add(name)
ifos.path.exists(orig_name):
self.path_set.existing_outputs.add(name)
returnposixpath.join(self.remote_dir,
self.remote_output_prefix,
name)

returnorig_name

classArgumentProcessor(PrefixProcessor):
def__init__(self,
input_prefix, output_prefix,
remote_dir, remote_input_prefix, remote_output_prefix,
path_set, arguments):
super().__init__(input_prefix, output_prefix,
remote_dir, remote_input_prefix, remote_output_prefix,
path_set)
self.original_args=arguments

defprocess_args(self):
self.args= []

forarginself.original_args:
self.args.append(self.process_one(arg))

classEnvVarProcessor(PrefixProcessor):
def__init__(self,
input_prefix, output_prefix,
remote_dir, remote_input_prefix, remote_output_prefix,
path_set, environment):
super().__init__(input_prefix, output_prefix,
remote_dir, remote_input_prefix, remote_output_prefix,
path_set)
self.original_env=environment

defprocess_env(self):
self.env=dict()

forkey, valueinself.original_env.items():
self.env[key] =self.process_one(value)

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', action='store_true',
help='run commands locally instead of over SSH, for '
'debugging purposes. The "host" argument is '
'omitted.')
parser.add_argument('--ignore-rsync-failure', action='store_true',
help='ignore rsync failures, for debugging.')

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.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
runner.ignore_rsync_failure=args.ignore_rsync_failure

assertnotargs.remote_dir=='/'

path_set=RemotePathSet()

argproc=ArgumentProcessor(args.input_prefix,
args.output_prefix,
args.remote_dir,
posixpath.normpath(args.remote_input_prefix),
posixpath.normpath(args.remote_output_prefix),
path_set,
args.command)

argproc.process_args()

envproc=EnvVarProcessor(args.input_prefix,
args.output_prefix,
args.remote_dir,
posixpath.normpath(args.remote_input_prefix),
posixpath.normpath(args.remote_output_prefix),
path_set,
collect_remote_env())

envproc.process_env()

input_dir=posixpath.join(args.remote_dir, args.remote_input_prefix)
output_dir=posixpath.join(args.remote_dir, args.remote_output_prefix)

ifargs.output_prefix:
runner.run_remote(['/bin/rm', '-rf', output_dir])

dirs=set()
foroutputinpath_set.outputs:
dirs.add(posixpath.join(output_dir, posixpath.dirname(output)))
foroutputinpath_set.nodir_outputs:
dirs.add(posixpath.join(output_dir, posixpath.dirname(output)))
runner.mkdirs_remote(list(dirs))

ifpath_set.inputs:
runner.send(args.input_prefix, input_dir, path_set.inputs)

ifpath_set.nodir_inputs:
runner.send(posixpath.dirname(args.input_prefix),
input_dir, path_set.nodir_inputs)

ifpath_set.existing_outputs:
runner.send(args.output_prefix, output_dir, path_set.existing_outputs)

ifpath_set.existing_nodir_outputs:
runner.send(posixpath.dirname(args.output_prefix),
output_dir, path_set.existing_nodir_outputs)

runner.run_remote(argproc.args, envproc.env)

ifpath_set.outputs:
runner.fetch(args.output_prefix, output_dir, path_set.outputs)

ifpath_set.nodir_outputs:
runner.fetch(posixpath.dirname(args.output_prefix),
output_dir, path_set.nodir_outputs)

if__name__=="__main__":
main()