utils/ssh.py

#===----------------------------------------------------------------------===##
#
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
#===----------------------------------------------------------------------===##

"""
Runs an executable on a remote host.

This is meant to be used as an executor when running the C++ Standard Library
conformance test suite.
"""

import argparse
import os
import subprocess
import sys


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', type=str, required=True)
    parser.add_argument('--codesign_identity', type=str, required=False)
    parser.add_argument('--dependencies', type=str, nargs='*', required=True)
    parser.add_argument('--env', type=str, nargs='*', required=True)
    (args, remaining) = parser.parse_known_args(sys.argv[1:])

    if len(remaining) < 2:
        sys.stderr.write('Missing actual commands to run')
        exit(1)
    remaining = remaining[1:] # Skip the '--'

    # HACK:
    # If the first argument is a file that ends in `.tmp.exe`, assume it is
    # the name of an executable generated by a test file. This allows us to
    # do custom processing like codesigning the executable and changing its
    # path when running on the remote host. It's possible for there to be no
    # such executable, for example in the case of a .sh.cpp test.
    exe = None
    if os.path.exists(remaining[0]) and remaining[0].endswith('.tmp.exe'):
        exe = remaining.pop(0)

    # If there's an executable, do any necessary codesigning.
    if exe and args.codesign_identity:
        rc = subprocess.call(['xcrun', 'codesign', '-f', '-s', args.codesign_identity, exe], env={})
        if rc != 0:
            sys.stderr.write('Failed to codesign: {}'.format(exe))
            return rc

    ssh = lambda command: ['ssh', '-oBatchMode=yes', args.host, command]
    scp = lambda src, dst: ['scp', '-oBatchMode=yes', '-r', src, '{}:{}'.format(args.host, dst)]

    # Create a temporary directory where the test will be run
    tmp = subprocess.check_output(ssh('mktemp -d /tmp/libcxx.XXXXXXXXXX')).strip()

    # Ensure the test dependencies exist and scp them to the temporary directory.
    # Test dependencies can be either files or directories, so the `scp` command
    # needs to use `-r`.
    for dep in args.dependencies:
        if not os.path.exists(dep):
            sys.stderr.write('Missing file or directory {} marked as a dependency of a test'.format(dep))
            exit(1)
        subprocess.call(scp(dep, tmp))

    # If there's an executable, change its path to be in the temporary directory.
    # We know it has been copied to the remote host when we handled the test
    # dependencies above.
    if exe:
        exe = os.path.join(tmp, os.path.basename(exe))

    # If there's an executable, make sure it has 'execute' permissions on the
    # remote host. The host that compiled the executable might not have a notion
    # of 'executable' permissions.
    if exe:
        subprocess.call(ssh('chmod +x {}'.format(exe)))

    # Execute the command through SSH in the temporary directory, with the
    # correct environment.
    commands = [
        'cd {}'.format(tmp),
        'export {}'.format(' '.join(args.env)),
        ' '.join([exe] + remaining if exe else remaining)
    ]
    res = subprocess.call(ssh(' && '.join(commands)))

    # Remove the temporary directory when we're done.
    subprocess.call(ssh('rm -r {}'.format(tmp)))

    return res

if __name__ == '__main__':
    exit(main())