libcxx/utils/ssh.py

   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 """
  11 Runs an executable on a remote host.
  12
  13 This is meant to be used as an executor when running the C++ Standard Library
  14 conformance test suite.
  15 """
  16
  17 import argparse
  18 import os
  19 import posixpath
  20 import shlex
  21 import subprocess
  22 import sys
  23 import tarfile
  24 import tempfile
  25
  26 from shlex import quote as cmd_quote
  27
  28
  29 def 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
  36 def 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
  43 def 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
 141 if __name__ == "__main__":
 142     exit(main())