blob: a53fe507ac430660cdbe98956284d73614732eb3 [file] [log] [blame]
#!/usr/bin/env python
# remote-run - Runs a command on another machine, for testing -----*- python -*-
# This source file is part of the open source project
# Copyright (c) 2018 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
# See for license information
# See 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
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.values())
self.run_remote(['/bin/mkdir', '-p'] + dirs_to_make)
# Send the local files.
sftp_commands = ("-put {0} {1}".format(quote(local_file),
for local_file, remote_file
in local_to_remote_files.items())
def fetch(self, local_to_remote_files):
# Prepare the local directory structure.
dirs_to_make = self._dirnames(local_to_remote_files.keys())
mkdir_command = ['/bin/mkdir', '-p'] + dirs_to_make
if self.verbose:
print(' '.join(mkdir_command), file=sys.stderr)
if not self.dry_run:
# Fetch the remote files.
sftp_commands = ("-get {0} {1}".format(quote(remote_file),
for local_file, remote_file
in local_to_remote_files.items())
def run_remote(self, command, remote_env={}):
env_strings = ['{0}={1}'.format(k,v) for k,v in sorted(remote_env.items())]
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:
_, _ = remote_proc.communicate()
if remote_proc.returncode:
# FIXME: We may still want to fetch the output files to see what
# went wrong.
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:
_, _ = sftp_proc.communicate(concatenated_commands.encode('utf-8'))
if 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)
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,
# with explanatory help from
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') +
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')
help='arguments matching this prefix will be uploaded')
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 '
parser.add_argument('--remote-output-prefix', default='output',
help='output arguments use this prefix on the remote '
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 '
help='the host to connect to, in the form '
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)
runner = RemoteCommandRunner(,
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,
assert not any(f in upload_files for f in 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,
assert not any(f in upload_files for f in test_files)
assert not any(f in download_files for f in 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:
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:
if __name__ == "__main__":