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 LLVM and the Libraries
14tests on a target.
15"""
16
17import argparse
18import os
19import posixpath
20import shlex
21import subprocess
22import sys
23import tarfile
24import tempfile
25import re
26
27def ssh(args, command):
28    cmd = ['ssh', '-oBatchMode=yes']
29    if args.extra_ssh_args is not None:
30        cmd.extend(shlex.split(args.extra_ssh_args))
31    return cmd + [args.host, command]
32
33
34def scp(args, src, dst):
35    cmd = ['scp', '-q', '-oBatchMode=yes']
36    if args.extra_scp_args is not None:
37        cmd.extend(shlex.split(args.extra_scp_args))
38    return cmd + [src, '{}:{}'.format(args.host, dst)]
39
40
41def main():
42    parser = argparse.ArgumentParser()
43    parser.add_argument('--host', type=str, required=True)
44    parser.add_argument('--execdir', type=str, required=False)
45    parser.add_argument('--extra-ssh-args', type=str, required=False)
46    parser.add_argument('--extra-scp-args', type=str, required=False)
47    parser.add_argument('--codesign_identity', type=str, required=False, default=None)
48    parser.add_argument('--env', type=str, nargs='*', required=False, default=dict())
49
50    # Note: The default value is for the backward compatibility with a hack in
51    # libcxx test suite.
52    # If an argument is a file that ends in `.tmp.exe`, assume it is the name
53    # of an executable generated by a test file. We call these test-executables
54    # below. This allows us to do custom processing like codesigning test-executables
55    # and changing their path when running on the remote host. It's also possible
56    # for there to be no such executable, for example in the case of a .sh.cpp
57    # test.
58    parser.add_argument('--exec-pattern', type=str, required=False, default='.*',
59                        help='The name regex pattern of the executables generated by \
60                              a test file. Specifying it allows us to do custom \
61                              processing like codesigning test-executables \
62                              and changing their path when running on \
63                              the remote host. It\'s also possible for there \
64                              to be no such executable, for example in \
65                              the case of a .sh.cpp test.')
66
67    parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
68    args = parser.parse_args()
69    commandLine = args.command
70
71    execdir = args.execdir
72    if execdir == '.':
73        # Retrieve the exec directory from the command line.
74        execdir, _ = os.path.split(commandLine[0])
75        if execdir == '':
76            # Get the current directory in that case.
77            execdir = os.getcwd()
78    arcname = os.path.basename(execdir) if execdir else None
79
80    # Create a temporary directory where the test will be run.
81    # That is effectively the value of %T on the remote host.
82    tmp = subprocess.check_output(
83                  ssh(args, 'mktemp -d /tmp/llvm.XXXXXXXXXX'),
84                  universal_newlines=True
85              ).strip()
86
87    isExecutable = lambda exe: re.match(args.exec_pattern, exe) and os.path.exists(exe)
88    pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))
89
90    remoteCommands = []
91
92    try:
93        # Do any necessary codesigning of test-executables found in the command line.
94        if args.codesign_identity:
95            for exe in filter(isExecutable, commandLine):
96                subprocess.check_call(
97                    ['xcrun', 'codesign', '-f', '-s', args.codesign_identity, exe],
98                    env={})
99
100        # tar up the execution directory (which contains everything that's needed
101        # to run the test), and copy the tarball over to the remote host.
102        if execdir:
103            try:
104                tmpTar = tempfile.NamedTemporaryFile(suffix='.tar', delete=False)
105                with tarfile.open(fileobj=tmpTar, mode='w') as tarball:
106                    tarball.add(execdir, arcname=arcname)
107
108                # Make sure we close the file before we scp it, because accessing
109                # the temporary file while still open doesn't work on Windows.
110                tmpTar.close()
111                remoteTarball = pathOnRemote(tmpTar.name)
112                subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
113            finally:
114                # Make sure we close the file in case an exception happens before
115                # we've closed it above -- otherwise close() is idempotent.
116                tmpTar.close()
117                os.remove(tmpTar.name)
118
119            # Untar the dependencies in the temporary directory and remove the tarball.
120            remoteCommands.extend([
121                'tar -xf {} -C {} --strip-components 1'.format(remoteTarball, tmp),
122                'rm {}'.format(remoteTarball)
123            ])
124        else:
125            # Copy only the files, which are specified in the command line.
126            # Copy them to remote host one by one.
127            for x in commandLine:
128                _, f = os.path.split(x)
129                subprocess.check_call(scp(args, x, pathOnRemote(f)))
130
131        # Make sure all executables in the remote command line have 'execute'
132        # permissions on the remote host. The host that compiled the test-executable
133        # might not have a notion of 'executable' permissions.
134        for exe in filter(isExecutable, commandLine):
135            remoteCommands.append('chmod +x {}'.format(pathOnRemote(exe)))
136
137        # Execute the command through SSH in the temporary directory, with the
138        # correct environment. We tweak the command line to run it on the remote
139        # host by transforming the path of test-executables to their path in the
140        # temporary directory on the remote host.
141        for i, x in enumerate(commandLine):
142            if isExecutable(x):
143                commandLine[i] = pathOnRemote(x)
144        remoteCommands.append('cd {}'.format(tmp))
145        if args.env:
146            remoteCommands.append('export {}'.format(' '.join(args.env)))
147        remoteCommands.append(subprocess.list2cmdline(commandLine))
148
149        # Finally, SSH to the remote host and execute all the commands.
150        rc = subprocess.call(ssh(args, ' && '.join(remoteCommands)))
151        return rc
152
153    finally:
154        # Make sure the temporary directory is removed when we're done.
155        subprocess.check_call(ssh(args, 'rm -r {}'.format(tmp)))
156
157
158if __name__ == '__main__':
159    exit(main())
160