#!/usr/local/munki/python
# encoding: utf-8
#
# Copyright 2011-2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
ptyexec

Created by John Randolph 2011-08-11.

Utility script to run a subprocess in a pseudo tty.

This will have the effect of unbuffering output I/O from the subprocess.
stdin of the subprocess is not connected to the stdin of this parent
process.
"""
from __future__ import absolute_import, print_function

import fcntl
import os
import pty
import select
import signal
import sys


# lots of variable names lifted from the C code this has been mostly
# translated from
# pylint: disable=invalid-name

child_exited = {}


def set_file_nonblocking(f, non_blocking=True):
    """Set non-blocking flag on a file object.

    Args:
        f: file
        non_blocking: bool, default True, non-blocking mode or not
    """
    flags = fcntl.fcntl(f.fileno(), fcntl.F_GETFL)
    if bool(flags & os.O_NONBLOCK) != non_blocking:
        flags ^= os.O_NONBLOCK
    fcntl.fcntl(f.fileno(), fcntl.F_SETFL, flags)


def sighandler(signum, _frame):
    """Handle a signal.

    Args:
        signum: int, signal number
        frame: frame, stack frame where signal was received
    """
    global child_exited

    if signum == signal.SIGCHLD:
        x = os.waitpid(-1, 0)
        if x[0] > 0:
            child_exited[x[0]] = x[1] >> 8    # get exit status from LSB


def usage(arg0):
    """Print usage."""
    print('Usage: %s [command to run] [arguments...]' % arg0, file=sys.stderr)
    return 0


def ptyexec(argv):
    """Setup pty and exec argv.

    Args:
        argv: list, arguments
    Returns:
        int, status code from child process upon completion, or 1 if a fork
            error occurs.
        never returns on the child side of the fork.
    """
    signal.signal(signal.SIGCHLD, sighandler)
    pid, fd = pty.fork()

    if pid == 0:         # child
        try:
            os.execv(argv[0], argv)
        except OSError as err:
            print(str(err), file=sys.stderr)
            sys.exit(1)
    elif pid > 0:        # parent
        f = os.fdopen(fd, 'rb+', 0)
        set_file_nonblocking(f)
        while True:
            try:
                (rl, _wl, _xl) = select.select([f], [], [], 5.0)
            except select.error:
                rl = []

            if f in rl:
                l = f.read()
                try:
                    sys.stdout.buffer.write(l)
                except AttributeError:
                    # Python 2
                    sys.stdout.write(l)
                sys.stdout.flush()

            if pid in child_exited:
                break

        f.close()
    elif pid == -1:    # error, never forked.
        print('fork() error', file=sys.stderr)
        return 1

    return child_exited.get(pid, 1)


def main(argv):
    """Main."""
    if len(argv) < 2:
        return usage(argv[0])

    argv.pop(0)
    return ptyexec(argv)


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