#!/usr/bin/env python
# ===----------------------------------------------------------------------===##
#
# 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 LLVM and the Libraries
tests on a target.
"""

import argparse
import os
import posixpath
import shlex
import subprocess
import sys
import tarfile
import tempfile
import re


def ssh(args, command):
    cmd = ["ssh", "-oBatchMode=yes"]
    if args.extra_ssh_args is not None:
        cmd.extend(shlex.split(args.extra_ssh_args))
    return cmd + [args.host, command]


def scp(args, src, dst):
    cmd = ["scp", "-q", "-oBatchMode=yes"]
    if args.extra_scp_args is not None:
        cmd.extend(shlex.split(args.extra_scp_args))
    return cmd + [src, "{}:{}".format(args.host, dst)]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, required=True)
    parser.add_argument("--execdir", type=str, required=False)
    parser.add_argument("--extra-ssh-args", type=str, required=False)
    parser.add_argument("--extra-scp-args", type=str, required=False)
    parser.add_argument("--codesign_identity", type=str, required=False, default=None)
    parser.add_argument("--env", type=str, nargs="*", required=False, default=dict())

    # Note: The default value is for the backward compatibility with a hack in
    # libcxx test suite.
    # If an argument is a file that ends in `.tmp.exe`, assume it is the name
    # of an executable generated by a test file. We call these test-executables
    # below. This allows us to do custom processing like codesigning test-executables
    # and changing their path when running on the remote host. It's also possible
    # for there to be no such executable, for example in the case of a .sh.cpp
    # test.
    parser.add_argument(
        "--exec-pattern",
        type=str,
        required=False,
        default=".*",
        help="The name regex pattern of the executables generated by \
                              a test file. Specifying it allows us to do custom \
                              processing like codesigning test-executables \
                              and changing their path when running on \
                              the remote host. It's also possible for there \
                              to be no such executable, for example in \
                              the case of a .sh.cpp test.",
    )

    parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
    args = parser.parse_args()
    commandLine = args.command

    execdir = args.execdir
    if execdir == ".":
        # Retrieve the exec directory from the command line.
        execdir, _ = os.path.split(commandLine[0])
        if execdir == "":
            # Get the current directory in that case.
            execdir = os.getcwd()
    arcname = os.path.basename(execdir) if execdir else None

    # Create a temporary directory where the test will be run.
    # That is effectively the value of %T on the remote host.
    tmp = subprocess.check_output(
        ssh(args, "mktemp -d /tmp/llvm.XXXXXXXXXX"), universal_newlines=True
    ).strip()

    isExecutable = lambda exe: re.match(args.exec_pattern, exe) and os.path.exists(exe)
    pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))

    remoteCommands = []

    try:
        # Do any necessary codesigning of test-executables found in the command line.
        if args.codesign_identity:
            for exe in filter(isExecutable, commandLine):
                subprocess.check_call(
                    ["xcrun", "codesign", "-f", "-s", args.codesign_identity, exe],
                    env={},
                )

        # tar up the execution directory (which contains everything that's needed
        # to run the test), and copy the tarball over to the remote host.
        if execdir:
            try:
                tmpTar = tempfile.NamedTemporaryFile(suffix=".tar", delete=False)
                with tarfile.open(fileobj=tmpTar, mode="w") as tarball:
                    tarball.add(execdir, arcname=arcname)

                # Make sure we close the file before we scp it, because accessing
                # the temporary file while still open doesn't work on Windows.
                tmpTar.close()
                remoteTarball = pathOnRemote(tmpTar.name)
                subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
            finally:
                # Make sure we close the file in case an exception happens before
                # we've closed it above -- otherwise close() is idempotent.
                tmpTar.close()
                os.remove(tmpTar.name)

            # Untar the dependencies in the temporary directory and remove the tarball.
            remoteCommands.extend(
                [
                    "tar -xf {} -C {} --strip-components 1".format(remoteTarball, tmp),
                    "rm {}".format(remoteTarball),
                ]
            )
        else:
            # Copy only files which are specified in the command line and exist on
            # the host. Copy them to the remote host one by one.
            for x in commandLine:
                if os.path.exists(x):
                    _, f = os.path.split(x)
                    subprocess.check_call(scp(args, x, pathOnRemote(f)))

        # Make sure all executables in the remote command line have 'execute'
        # permissions on the remote host. The host that compiled the test-executable
        # might not have a notion of 'executable' permissions.
        for exe in filter(isExecutable, commandLine):
            remoteCommands.append("chmod +x {}".format(pathOnRemote(exe)))

        # Execute the command through SSH in the temporary directory, with the
        # correct environment. We tweak the command line to run it on the remote
        # host by transforming the path of test-executables to their path in the
        # temporary directory on the remote host.
        for i, x in enumerate(commandLine):
            if isExecutable(x):
                commandLine[i] = pathOnRemote(x)
        remoteCommands.append("cd {}".format(tmp))
        if args.env:
            remoteCommands.append("export {}".format(" ".join(args.env)))
        remoteCommands.append(subprocess.list2cmdline(commandLine))

        # Finally, SSH to the remote host and execute all the commands.
        rc = subprocess.call(ssh(args, " && ".join(remoteCommands)))
        return rc

    finally:
        # Make sure the temporary directory is removed when we're done.
        subprocess.check_call(ssh(args, "rm -r {}".format(tmp)))


if __name__ == "__main__":
    rc = main()

    # If the remote process died with a signal, this will be exposed by ssh as
    # an exit code in this range. We may be running under `not --crash` which
    # will expect us to also die with a signal, so send the signal to ourselves
    # so that wait4() in `not` will detect the signal. 
    if rc > 128 and rc < 160:
        os.kill(os.getpid(), rc - 128)

    exit(rc)