#!/usr/bin/python
#
# Copyright (c) 2018 Red Hat, Inc.
#
# This software is licensed to you under the GNU General Public License,
# version 3 (GPLv3). There is NO WARRANTY for this software, express or
# implied, including the implied warranties of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. You should have received a copy of GPLv3
# along with this software; if not, see
# https://www.gnu.org/licenses/gpl-3.0.txt.
"""Run ssh with a timeout."""

import os
import platform
import sys

ARG_EXECUTABLE = '--executable='
ARG_TIMEOUT = '--timeout='
ARG_SSH = 'ssh'

# Ansible lets us set the ssh executable, but whatever string we
# provide is interpreted as an executable, so we can't set it to
# 'timeout 100' or something like that. Instead, we have this tiny
# program whose only job is to be a single executable that Ansible can
# call which will then trampoline out to timeout.


def main(argv):
    # I don't think we can control the order that Ansible passes
    # arguments to its ssh program, so we can't just make hte
    # executable the first positional argument. Instead, we pass the
    # executable as --executable=<program>, and then we remove that
    # flag before we pass the rest of the args to the executable.
    executable = None
    timeout = None
    for arg in argv[1:]:
        if arg.startswith(ARG_EXECUTABLE):
            executable = arg[len(ARG_EXECUTABLE):]
        # timeout works the same way as executable
        if arg.startswith(ARG_TIMEOUT):
            timeout = arg[len(ARG_TIMEOUT):]

    if not executable:
        print('Error: must pass --executable to timeout_ssh. Got:', argv)
        sys.exit(1)

    if not timeout:
        print('Error: must pass --timeout to timeout_ssh. Got:', argv)
        sys.exit(1)

    if platform.system() == 'Darwin':
        timeout_command = 'gtimeout'
    else:
        timeout_command = 'timeout'

    # the ssh arg is required for become-pass because
    # ansible checks for an exact string match of ssh
    # anywhere in the command array
    # See https://github.com/ansible/ansible/blob/stable-2.3/lib/ansible/plugins/connection/ssh.py#L490-L500
    # the code below will remove the ssh argument before running the command
    executable_args = [arg for arg in argv[1:]
                       if (not arg.startswith(ARG_EXECUTABLE) and
                           not arg.startswith(ARG_TIMEOUT) and
                           arg != ARG_SSH)]

    timeout_args = [timeout_command, timeout, executable] + executable_args

    os.execvp(timeout_command, timeout_args)


if __name__ == '__main__':
    main(sys.argv)
