246 строки
10 KiB
Python
Executable File
246 строки
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.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()
|