 ############################################################################
 #                                                                          #
 #                              EX.PY                                       #
 #                                                                          #
 #           Copyright (C) 2008 - 2011 Ada Core Technologies, Inc.          #
 #                                                                          #
 # This program is free software: you can redistribute it and/or modify     #
 # it under the terms of the GNU General Public License as published by     #
 # the Free Software Foundation, either version 3 of the License, or        #
 # (at your option) any later version.                                      #
 #                                                                          #
 # This program is distributed in the hope that it will be useful,          #
 # but WITHOUT ANY WARRANTY; without even the implied warranty of           #
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
 # GNU General Public License for more details.                             #
 #                                                                          #
 # You should have received a copy of the GNU General Public License        #
 # along with this program.  If not, see <http://www.gnu.org/licenses/>     #
 #                                                                          #
 ############################################################################

"""Subprocesses management

This package provides a single class called run which ease spawn of processes
in blocking or non blocking mode and redirection of its stdout, stderr and
stdin"""

from subprocess import Popen, STDOUT, PIPE

import errno
import logging
import os
import sys

BUF_SIZE = 128

logger = logging.getLogger('gnatpython.ex')


class Run(object):
    """
    ATTRIBUTES
      cmds   : The `cmds' argument passed to the __init__ method
               (a command line passed in a list, of a list of command
               lines passed as a list of list).
      status : exit status (meaningfull only after the end of the process)
      out    : process standard output  (if instanciated with output = PIPE)
      err    : same as out but for standard error
      pid    : PID
    """

    def __init__(self, cmds, cwd=None, output=PIPE,
                 error=STDOUT, input=None, bg=False, timeout=None, env=None,
                 set_sigpipe=True, parse_shebang=False):
        """Spawn a process

        PARAMETERS
          cmds:    two possibilities:
                   (1) a command line: a tool name and its arguments, passed
                   in a list. e.g. ['ls', '-a', '.']
                   (2) a list of command lines (as defined in (1)): the
                   different commands will be piped. This means that
                   [['ps', '-a'], ['grep', 'vxsim']] will be equivalent to
                   the system command line 'ps -a | grep vxsim'.
          cwd :    directory in which the process should be executed (string
                   or None). If None then current directory is used
          output:  can be PIPE (default), a filename string, a fd on an already
                   opened file, a python file object or None (for stdout).
          error:   same as output or STDOUT, which indicates that the stderr
                   data from the applications should be captured into the
                   same file handle as for stdout.
          input:   same as output
          bg:      if True then run in background
          timeout: limit execution time (in seconds)
          env:     dictionary for environment variables (e.g. os.environ)
          set_sigpipe: reset SIGPIPE handler to default value
          parse_shebang: take the #! interpreter line into account

        RETURN VALUE
          Return an object of type run.

        EXCEPTIONS
          Raise OSError when trying to execute a non-existent file.

        REMARKS
          If you specify a filename for output or stderr then file content is
          reseted (equiv. to > in shell). If you prepend the filename with '+'
          then the file will be opened in append mode (equiv. to >> in shell)
          If you prepend the input with '|', then the content of input string
          will be used for process stdin.
        """

        def subprocess_setup():
            """Reset SIGPIPE hander

            Python installs a SIGPIPE handler by default. This is usually not
            what non-Python subprocesses expect.
            """
            if set_sigpipe:
                # Set sigpipe only when set_sigpipe is True
                # This should fix HC16-020 and could be activated by default
                import signal
                signal.signal(signal.SIGPIPE, signal.SIG_DFL)

        def add_interpreter_command(cmd_line):
            """Add the interpreter defined in the #! line to cmd_line

            If the #! line cannot be parsed, just return the cmd_line
            unchanged

            REMARKS
              if the interpreter command line contains /usr/bin/env python
              it will be replaced by the value of sys.executable

              On windows, /usr/bin/env will be ignored to avoid a dependency on
              cygwin
            """
            if not parse_shebang:
                # nothing to do
                return cmd_line
            # Import gnatpython.fileutils just now to avoid a circular
            # dependency
            from gnatpython.fileutils import which
            prog = which(cmd_line[0])
            if not os.path.exists(prog):
                return cmd_line
            with open(prog) as f:
                header = f.read()[0:2]
                if header != "#!":
                    # Unknown header
                    return cmd_line
                # Header found, get the interpreter command in the first line
                f.seek(0)
                line = f.readline()
                interpreter_cmds = [l.strip() for l in
                        line[line.find('!') + 1:].split()]
                # Pass the program path to the interpreter
                if len(cmd_line) > 1:
                    cmd_line = [prog] + cmd_line[1:]
                else:
                    cmd_line = [prog]

                # If the interpreter is '/usr/bin/env python', use
                # sys.executable instead to keep the same python executable
                if interpreter_cmds[0:1] == ['/usr/bin/env', 'python']:
                    if len(interpreter_cmds > 2):
                        return [sys.executable] + interpreter_cmds[2:] \
                                + cmd_line
                    else:
                        return [sys.executable] + cmd_line
                elif sys.platform == 'win32':
                    if interpreter_cmds[0] == '/usr/bin/env':
                        return interpreter_cmds[1:] + cmd_line
                return interpreter_cmds + cmd_line

        # First resolve output, error and input
        self.input_file = File(input, 'r')
        self.output_file = File(output, 'w')
        self.error_file = File(error, 'w')

        self.status = None
        self.out = ''
        self.err = ''

        if env is None:
            env = os.environ

        rlimit_args = []
        if timeout is not None:
            # Import gnatpython.fileutils just now to avoid a circular
            # dependency
            from gnatpython.fileutils import get_rlimit
            rlimit = get_rlimit()
            assert rlimit, 'rlimit not found'
            rlimit_args = [rlimit, '%d' % timeout]

        try:
            if not isinstance(cmds[0], list):
                self.cmds = rlimit_args + add_interpreter_command(cmds)
                logger.debug('Run: %s' % self.command_line_image())

                popen_args = {
                    'stdin': self.input_file.fd,
                    'stdout': self.output_file.fd,
                    'stderr': self.error_file.fd,
                    'cwd': cwd,
                    'env': env,
                    'universal_newlines': True}

                if sys.platform != 'win32':
                    # preexec_fn is no supported on windows
                    popen_args['preexec_fn'] = subprocess_setup

                self.internal = Popen(self.cmds, **popen_args)

            else:
                self.cmds = [add_interpreter_command(c) for c in cmds]
                self.cmds[0] = rlimit_args + self.cmds[0]

                logger.debug('Run: %s ' %
                              " | ".join([" ".join(cmd) for cmd in self.cmds]))
                runs = []
                for index, cmd in enumerate(self.cmds):
                    if index == 0:
                        stdin = self.input_file.fd
                    else:
                        stdin = runs[index - 1].stdout

                    # When connecting two processes using a Pipe don't use
                    # universal_newlines mode. Indeed commands transmitting
                    # binary data between them will crash
                    # (ex: gzip -dc toto.txt | tar -xf -)
                    if index == len(self.cmds) - 1:
                        stdout = self.output_file.fd
                        txt_mode = True
                    else:
                        stdout = PIPE
                        txt_mode = False

                    popen_args = {
                        'stdin': stdin,
                        'stdout': stdout,
                        'stderr': self.error_file.fd,
                        'cwd': cwd,
                        'env': env,
                        'universal_newlines': txt_mode}

                    if sys.platform != 'win32':
                        # preexec_fn is no supported on windows
                        popen_args['preexec_fn'] = subprocess_setup

                    runs.append(Popen(cmd, **popen_args))
                    self.internal = runs[-1]

        except Exception, e:
            self.__error(e, self.cmds)
            raise

        self.pid = self.internal.pid

        if not bg:
            self.wait()

    def command_line_image(self):
        """Return a string representation of the command(s) that
        were run to create this object.

        REMARKS
            This method also handles quoting as defined for POSIX shells.
            This means that arguments containing special characters
            (such as a simple space, or a backslash, for instance),
            are properly quoted.  This makes it possible to execute
            the same command by copy/pasting the image in a shell
            prompt.
        """
        def quote_arg(arg):
            """Return a human-friendly representation of the given
            argument, but with all extra quoting done if necessary.
            The intent is to produce an argument image that can be
            copy/pasted on a POSIX shell command (at a shell prompt).
            """
            need_quoting = ('|', '&', ';', '<',  '>',  '(', ')',  '$',
                            '`', '\\',  '"',  "'", ' ', '\t', '\n',
                            # The POSIX spec says that the following
                            # characters might need some extra quoting
                            # depending on the circumstances.  We just
                            # always quote them, to be safe (and to avoid
                            # things like file globbing which are sometimes
                            # performed by the shell). We do leave '%' and
                            # '=' alone, as I don't see how they could
                            # cause problems.
                            '*', '?', '[', '#', '~')
            for char in need_quoting:
                if char in arg:
                    # The way we do this is by simply enclosing the argument
                    # inside single quotes.  However, we have to be careful
                    # of single-quotes inside the argument, as they need
                    # to be escaped (which we cannot do while still inside.
                    # a single-quote string).
                    arg = arg.replace("'", r"'\''")
                    # Also, it seems to be nicer to print new-line characters
                    # as '\n' rather than as a new-line...
                    arg = arg.replace('\n', r"'\n'")
                    return "'%s'" % arg
            # No quoting needed.  Return the argument as is.
            return arg

        cmds = self.cmds
        if not isinstance(cmds[0], list):
            # Turn the simple command into a special case of
            # the multiple-commands case.  This will allow us
            # to treat both cases the same way.
            cmds = [cmds]
        return ' | '.join([' '.join([quote_arg(arg) for arg in cmd])
                           for cmd in cmds])

    def _close_files(self):
        """Internal procedure"""
        self.output_file.close()
        self.error_file.close()
        self.input_file.close()

    def __error(self, error, cmds):
        """Set pid to -1 and status to 127 before closing files"""
        self.pid = -1
        self.status = 127
        self._close_files()

        # Try to send an helpful message if one of the executable has not
        # been found.
        not_found = None

        # Import gnatpython.fileutils here to avoid a circular dependency
        from gnatpython.fileutils import which
        if not isinstance(cmds[0], list):
            if not which(cmds[0]):
                not_found = cmds[0]
        else:
            for cmd in cmds:
                if not which(cmd[0]):
                    not_found = cmd[0]
                    break

        if not_found is not None:
            logger.error("%s, %s not found" % (error, not_found))
            raise OSError(getattr(error, 'errno', errno.ENOENT),
                    getattr(error, 'strerror', 'No such file or directory') +
                    " %s not found" % not_found)

    def wait(self):
        """Wait until process ends and return its status"""
        if self.status == 127:
            return self.status

        self.status = None

        # If there is no pipe in the loop then just do a wait. Otherwise
        # in order to avoid blocked processes due to full pipes, use
        # communicate.
        if self.output_file.fd != PIPE and self.error_file.fd != PIPE and \
          self.input_file.fd != PIPE:
            self.status = self.internal.wait()
        else:
            tmp_input = None
            if self.input_file.fd == PIPE:
                tmp_input = self.input_file.get_command()

            (self.out, self.err) = self.internal.communicate(tmp_input)
            self.status = self.internal.returncode

        self._close_files()
        return self.status

    def poll(self):
        """Test if the process is still alive. If yes then return None,
        otherwise return process status"""

        if self.status != 127:
            result = self.internal.poll()
            if result is not None:
                self.status = result
        else:
            result = 127
        return result


class File(object):
    """Can be a PIPE, a file object"""

    def __init__(self, name, mode='r'):
        """Create a new File

        PARAMETERS
          name: can be PIPE, STDOUT, a filename string,
                an opened fd, a python file object,
                or a command to pipe (if starts with |)

          mode: can be 'r' or 'w'
                if name starts with + the mode will be a+
        """
        assert mode in 'rw', 'Mode should be r or w'

        self.name = name
        self.to_close = False
        if isinstance(name, str):
            # can be a pipe or a filename
            if mode == 'r' and name.startswith('|'):
                self.fd = PIPE
            else:
                if mode == 'w':
                    if name.startswith('+'):
                        open_mode = 'a+'
                        name = name[1:]
                    else:
                        open_mode = 'w+'
                else:
                    open_mode = 'r'

                self.fd = open(name, open_mode)
                if open_mode == 'a+':
                    self.fd.seek(0, 2)
                self.to_close = True

        else:
            # this is a file descriptor
            self.fd = name

    def get_command(self):
        """Returns the command to run to create the pipe"""
        if self.fd == PIPE:
            return self.name[1:]

    def close(self):
        """Close the file if needed"""
        if self.to_close:
            self.fd.close()
