xref: /llvm-project/libcxx/utils/ssh.py (revision 7bfaa0f09d0564f315ea778023b34b8a113ec740)
1#!/usr/bin/env python
2# ===----------------------------------------------------------------------===##
3#
4# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5# See https://llvm.org/LICENSE.txt for license information.
6# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7#
8# ===----------------------------------------------------------------------===##
9
10"""
11Runs an executable on a remote host.
12
13This is meant to be used as an executor when running the C++ Standard Library
14conformance test suite.
15"""
16
17import argparse
18import os
19import posixpath
20import shlex
21import subprocess
22import sys
23import tarfile
24import tempfile
25
26from shlex import quote as cmd_quote
27
28
29def ssh(args, command):
30    cmd = ["ssh", "-oBatchMode=yes"]
31    if args.extra_ssh_args is not None:
32        cmd.extend(shlex.split(args.extra_ssh_args))
33    return cmd + [args.host, command]
34
35
36def scp(args, src, dst):
37    cmd = ["scp", "-q", "-oBatchMode=yes"]
38    if args.extra_scp_args is not None:
39        cmd.extend(shlex.split(args.extra_scp_args))
40    return cmd + [src, "{}:{}".format(args.host, dst)]
41
42
43def main():
44    parser = argparse.ArgumentParser()
45    parser.add_argument("--host", type=str, required=True)
46    parser.add_argument("--execdir", type=str, required=True)
47    parser.add_argument("--tempdir", type=str, required=False, default="/tmp")
48    parser.add_argument("--extra-ssh-args", type=str, required=False)
49    parser.add_argument("--extra-scp-args", type=str, required=False)
50    parser.add_argument("--codesign_identity", type=str, required=False, default=None)
51    parser.add_argument("--env", type=str, nargs="*", required=False, default=[])
52    parser.add_argument(
53        "--prepend_env", type=str, nargs="*", required=False, default=[]
54    )
55    parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
56    args = parser.parse_args()
57    commandLine = args.command
58
59    # Create a temporary directory where the test will be run.
60    # That is effectively the value of %T on the remote host.
61    tmp = subprocess.check_output(
62        ssh(args, "mktemp -d {}/libcxx.XXXXXXXXXX".format(args.tempdir)),
63        universal_newlines=True,
64    ).strip()
65
66    # HACK:
67    # If an argument is a file that ends in `.tmp.exe`, assume it is the name
68    # of an executable generated by a test file. We call these test-executables
69    # below. This allows us to do custom processing like codesigning test-executables
70    # and changing their path when running on the remote host. It's also possible
71    # for there to be no such executable, for example in the case of a .sh.cpp
72    # test.
73    isTestExe = lambda exe: exe.endswith(".tmp.exe") and os.path.exists(exe)
74    pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))
75
76    try:
77        # Do any necessary codesigning of test-executables found in the command line.
78        if args.codesign_identity:
79            for exe in filter(isTestExe, commandLine):
80                subprocess.check_call(
81                    ["xcrun", "codesign", "-f", "-s", args.codesign_identity, exe],
82                    env={},
83                )
84
85        # tar up the execution directory (which contains everything that's needed
86        # to run the test), and copy the tarball over to the remote host.
87        try:
88            tmpTar = tempfile.NamedTemporaryFile(suffix=".tar", delete=False)
89            with tarfile.open(fileobj=tmpTar, mode="w") as tarball:
90                tarball.add(args.execdir, arcname=os.path.basename(args.execdir))
91
92            # Make sure we close the file before we scp it, because accessing
93            # the temporary file while still open doesn't work on Windows.
94            tmpTar.close()
95            remoteTarball = pathOnRemote(tmpTar.name)
96            subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
97        finally:
98            # Make sure we close the file in case an exception happens before
99            # we've closed it above -- otherwise close() is idempotent.
100            tmpTar.close()
101            os.remove(tmpTar.name)
102
103        # Untar the dependencies in the temporary directory and remove the tarball.
104        remoteCommands = [
105            "tar -xf {} -C {} --strip-components 1".format(remoteTarball, tmp),
106            "rm {}".format(remoteTarball),
107        ]
108
109        # Make sure all test-executables in the remote command line have 'execute'
110        # permissions on the remote host. The host that compiled the test-executable
111        # might not have a notion of 'executable' permissions.
112        for exe in map(pathOnRemote, filter(isTestExe, commandLine)):
113            remoteCommands.append("chmod +x {}".format(exe))
114
115        # Execute the command through SSH in the temporary directory, with the
116        # correct environment. We tweak the command line to run it on the remote
117        # host by transforming the path of test-executables to their path in the
118        # temporary directory on the remote host.
119        commandLine = (pathOnRemote(x) if isTestExe(x) else x for x in commandLine)
120        remoteCommands.append("cd {}".format(tmp))
121
122        if args.prepend_env:
123            # We can't sensibly know the original value of the env vars
124            # in order to prepend to them, so just overwrite these variables.
125            args.env.extend(args.prepend_env)
126
127        if args.env:
128            env = list(map(cmd_quote, args.env))
129            remoteCommands.append("export {}".format(" ".join(args.env)))
130        remoteCommands.append(subprocess.list2cmdline(commandLine))
131
132        # Finally, SSH to the remote host and execute all the commands.
133        rc = subprocess.call(ssh(args, " && ".join(remoteCommands)))
134        return rc
135
136    finally:
137        # Make sure the temporary directory is removed when we're done.
138        subprocess.check_call(ssh(args, "rm -r {}".format(tmp)))
139
140
141if __name__ == "__main__":
142    exit(main())
143