mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
2b8d57c2e3
With Python 3, there have been some sorting inconsistencies. Always sort the items to allow for the lit test to reliably match the output.
246 lines
10 KiB
Python
Executable File
246 lines
10 KiB
Python
Executable File
#!/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.values())
|
|
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.items())
|
|
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.keys())
|
|
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.items())
|
|
self.run_sftp(sftp_commands)
|
|
|
|
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:
|
|
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.encode('utf-8'))
|
|
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(f in upload_files 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(f in upload_files for f in test_files)
|
|
upload_files.update(test_files)
|
|
assert not any(f in download_files 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()
|