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
#
# ----------------------------------------------------------------------------

import argparse
import os
import posixpath
import subprocess
import sys
import shutil

def quote(arg):
    return repr(arg)

class CommandRunner(object):
    def __init__(self):
        self.verbose = False
        self.dry_run = False
        self.ignore_rsync_failure = 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 mkdirs_remote(self, directories):
        if directories:
            mkdir_command = ['/bin/mkdir', '-p'] + directories
            self.run_remote(mkdir_command)

    def send(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)

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

        self.run_rsync_from(remote_prefix, remote_to_local_files, output_prefix)

    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_rsync(self, invocation, sources):
        rsync_proc = self.popen(invocation,
                                stdin=subprocess.PIPE,
                                stdout=None, stderr=None)
        if self.dry_run:
            return
        sources = '\n'.join(sources)
        if self.verbose:
            print(sources, file=sys.stderr)
        _, _ = rsync_proc.communicate(sources.encode('utf-8'))
        if not self.ignore_rsync_failure and rsync_proc.returncode:
            sys.exit(rsync_proc.returncode)

    def run_rsync_to(self, prefix, sources, dest):
        self.run_rsync(self.rsync_to_invocation(prefix, dest), sources)

    def run_rsync_from(self, prefix, sources, dest):
        self.run_rsync(self.rsync_from_invocation(prefix, dest), sources)

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 rsync_invocation(self, source, dest):
        return ['/usr/bin/rsync', '-arRz', '--files-from=-',
                '-e', ' '.join([quote(x) for x in
                                ['/usr/bin/ssh'] +
                                self.common_options(port_flag='-p')]),
                source,
                dest]

    def rsync_from_invocation(self, source, dest):
        return self.rsync_invocation(self.remote_host + ':' + source, dest)

    def rsync_to_invocation(self, source, dest):
        return self.rsync_invocation(source, self.remote_host + ':' + dest)

class LocalCommandRunner(CommandRunner):
    def remote_invocation(self, command):
        return command

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

    def rsync_from_invocation(self, source, dest):
        return self.rsync_invocation(source, dest)

    def rsync_to_invocation(self, source, dest):
        return self.rsync_invocation(source, dest)

def strip_sep(name):
    lsep = len(posixpath.sep)
    while name.startswith(posixpath.sep):
        name = name[lsep:]
    return name

class ArgumentProcessor(object):
    def __init__(self,
                 input_prefix, output_prefix,
                 remote_dir, remote_input_prefix, remote_output_prefix,
                 arguments):
        assert not remote_input_prefix.startswith('..')
        assert not remote_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.original_args = arguments

        if self.input_prefix:
            while self.input_prefix.endswith(posixpath.sep):
                self.input_prefix = self.input_prefix[:-len(posixpath.sep)]
        if self.output_prefix:
            while self.output_prefix.endswith(posixpath.sep):
                self.output_prefix = self.output_prefix[:-len(posixpath.sep)]

    def process_args(self):
        self.args = []
        self.inputs = set()
        self.nodir_inputs = set()
        self.existing_outputs = set()
        self.outputs = set()
        self.existing_nodir_outputs = set()
        self.nodir_outputs = set()

        iplen = 0
        oplen = 0
        if self.input_prefix:
            iplen = len(self.input_prefix)
            ipfxdir, ipfx = posixpath.split(self.input_prefix)
        if self.output_prefix:
            oplen = len(self.output_prefix)
            opfxdir, opfx = posixpath.split(self.output_prefix)
        for arg in self.original_args:
            if iplen and arg.startswith(self.input_prefix):
                name = arg[iplen:]
                if not name.startswith(posixpath.sep):
                    name = ipfx + name
                    self.nodir_inputs.add(name)
                else:
                    name = strip_sep(name)
                    self.inputs.add(name)

                self.args.append(posixpath.join(self.remote_dir,
                                                self.remote_input_prefix,
                                                name))
            elif oplen and arg.startswith(self.output_prefix):
                name = arg[oplen:]
                if not name.startswith(posixpath.sep):
                    name = opfx + name
                    self.nodir_outputs.add(name)
                    if os.path.exists(arg):
                        self.existing_nodir_outputs.add(name)
                else:
                    name = strip_sep(name)
                    self.outputs.add(name)
                    if os.path.exists(arg):
                        self.existing_outputs.add(name)

                self.args.append(posixpath.join(self.remote_dir,
                                                self.remote_output_prefix,
                                                name))
            else:
                self.args.append(arg)

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', 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()

    if args.debug_as_local:
        runner = LocalCommandRunner()
        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
    runner.ignore_rsync_failure = args.ignore_rsync_failure

    assert not args.remote_dir == '/'

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

    argproc.process_args()

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

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

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

    if argproc.inputs:
        runner.send(args.input_prefix, input_dir, argproc.inputs)

    if argproc.nodir_inputs:
        runner.send(posixpath.dirname(args.input_prefix),
                    input_dir, argproc.nodir_inputs)

    if argproc.existing_outputs:
        runner.send(args.output_prefix, output_dir, argproc.existing_outputs)

    if argproc.existing_nodir_outputs:
        runner.send(posixpath.dirname(args.output_prefix),
                    output_dir, argproc.existing_nodir_outputs)

    remote_env = collect_remote_env()

    runner.run_remote(argproc.args, remote_env)

    if argproc.outputs:
        runner.fetch(args.output_prefix, output_dir, argproc.outputs)

    if argproc.nodir_outputs:
        runner.fetch(posixpath.dirname(args.output_prefix),
                     output_dir, argproc.nodir_outputs)

if __name__ == "__main__":
    main()