#!/usr/bin/python

########################################################################
# Copyright 2009-2010,2016-2023 VMware, Inc.  All rights reserved.
# -- VMware Confidential
########################################################################

"""This module provides the runcommand function, a convenience function to
   run a command and collect the output.
"""

import errno
import fcntl
import logging
import os
import select
import signal
import subprocess
import sys
import time

log = logging.getLogger(__name__)

MAXOUTPUTSIZE = 1048576 # 1 MB
ONVISOR = os.uname()[0].lower().startswith('vmkernel')
try:
   import vmkuser
   HAVEVMKUSER = True
except:
   HAVEVMKUSER = False

class RunCommandError(Exception):
   pass

def isString(objInstance):
   """Check whether a given object is a string
   """
   if sys.version_info[0] >= 3:
      return isinstance(objInstance, str)
   else:
      return isinstance(objInstance, basestring)

class VisorPopen(object):
   """A rather simple-minded emulation of the Popen class for use on Visor."""

   def __init__(self, args, redirectErr=True, stdin=None):
      """Class constructor.
            Parameters:
               * args - Either a sequence or string. If a string, the string
                        will be executed via a shell. If a sequence, the first
                        item is the program to execute, and the remaining items
                        are arguments to the program.
               * redirectErr - If true, the error will be redirected to the output
               * stdin - stdin file descriptor the subprocess should use.
         Returns: A VisorPopen object.
         Raises:
            * RuntimeError - If the vmkuser module is not available.
            * OSError      - If the command cannot be executed.
      """
      self.returncode = None
      self.stdout = None
      self.stderr = None

      if not HAVEVMKUSER:
         msg = "Class %s requires vmkuser module" % self.__class__.__name__
         raise RuntimeError(msg)

      if isString(args):
         args = ['/bin/sh', '-c', args]
      else:
         args = list(args)

      fromchildfd, toparentfd = os.pipe()

      if not stdin:
         stdin = sys.stdin.fileno()

      if not redirectErr:
         fromchilderrfd, toparenterrfd = os.pipe()
         initfds = [stdin, toparentfd, toparenterrfd]
      else:
         initfds = [stdin, toparentfd, toparentfd]

      try:
         self.pid = vmkuser.ForkExec(args[0],  # filePath
                                     args,     # argv
                                     None,     # envp
                                     -1,       # workingDirFd
                                     initfds,  # initFds
                                     -1,       # uid
                                     -1,       # gid
                                     0)        # detached
      except:
         os.close(fromchildfd)
         if not redirectErr:
            os.close(fromchilderrfd)
         raise
      finally:
         os.close(toparentfd)
         if not redirectErr:
            os.close(toparenterrfd)

      # Open Python file objects only if ForkExec succeeds.
      self.stdout = os.fdopen(fromchildfd, 'rb')
      if not redirectErr:
         self.stderr = os.fdopen(fromchilderrfd, 'rb')

   def __del__(self):
      # Ensure the stdout/stderr files are closed to avoid a ResourceWarning.
      if self.stdout:
         self.stdout.close()
      if self.stderr:
         self.stderr.close()

   # The rest of this is shamelessly borrowed from Python's subprocess module.
   def _handle_exitstatus(self, sts):
      if os.WIFSIGNALED(sts):
         self.returncode = -os.WTERMSIG(sts)
      elif os.WIFEXITED(sts):
         self.returncode = os.WEXITSTATUS(sts)
      else:
         # Should never happen
         raise RuntimeError("Unknown child exit status!")

   def poll(self):
      if self.returncode == None:
         try:
            pid, sts = os.waitpid(self.pid, os.WNOHANG)
            if pid == self.pid:
               self._handle_exitstatus(sts)
         except os.error as e:
            if e.errno == errno.ECHILD:
               # This happens if SIGCHLD is set to be ignored or waiting
               # for child processes has otherwise been disabled for our
               # process. We cannot get the return status of the child
               # anymore.
               self.returncode = 0
      return self.returncode

def terminateprocess(p):
   # Get progressively heavy-handed.
   for s in (signal.SIGINT, signal.SIGTERM, signal.SIGKILL):
      try:
         os.kill(p.pid, s)
      except OSError:
         # Maybe the process already exited. We could check for ESRCH, but we
         # wouldn't really do anything differently, anyway.
         return
      timeout = 3.0
      while timeout > 0.0:
         # Use select as a per-thread 'sleep'.
         select.select([], [], [], 0.02)
         timeout -= 0.02
         if p.poll() != None:
            return

def terminateProcessGroup(
   proc,
   termSignals=(signal.SIGINT, signal.SIGTERM),
   waitTimeout=3,
):
   """Terminate process-group led by `proc`; wait for group to disappear.

   WARNING: Does NOT enforce `proc` is actually a process-group leader.

   Parameters:
      * proc - `subprocess.Popen` or `VisorPopen` instance.
      * termSignals - Termination signals. `SIGKILL` is always the last-resort.
      * waitTimeout - Seconds to wait before trying the next signal.
   Returns:
      * On success, nothing; `proc` will have its `returncode` etc. set.
      * Raises `RunCommandError` if the group does not disappear after all
      signals have been tried.
   """
   # By definition, the group-leader's process ID is the group's ID.
   pgid = proc.pid

   # Give the group a chance to clean up.
   # In the last two rounds, send a SIGKILL, followed by "signal" 0 as a final
   # group-existence check.
   for termSignal in (*termSignals, signal.SIGKILL, 0):
      log.warning('Attempting to kill process group %d '
                  '(termSignal=%r, waitTimeout=%d)',
                  pgid, termSignal, waitTimeout)
      try:
         os.killpg(pgid, termSignal)
      except ProcessLookupError:
         log.info('Process group %d already terminated', pgid)
         return   # Group no longer exists - done!

      if termSignal == 0:
         # Give up (raise error below).
         log.warning('Process group %d still active!', pgid)
         break

      # Wait at least once for `proc`'s termination (ensure reaping).
      # NOTE: `Popen.poll()` is a no-op if it previously succeeded.
      proc.poll()

      # The group may still exist since `killpg()` succeeded. Wait a bit.
      time.sleep(waitTimeout)

   raise RunCommandError('Process group {} still active!'.format(pgid))

def runcommand(args, outfile=None, returnoutput=True, timeout=0.0,
               redirectErr=True, maxOutputSize=MAXOUTPUTSIZE, stdinStr=None):
   """Execute a command, optionally collecting output.
         Parameters:
            * args         - May be either a sequence of arguments, or a
                             string. If a string, the string will be passed to
                             a shell for execution. If a sequence of arguments,
                             the first item is the program to execute, and the
                             remaining items are arguments to the program. Note
                             that when a sequence of arguments is specified,
                             the shell is not invoked, so arguments are not
                             subject to shell interpolation.
            * outfile      - If specified, may be either a file-like object
                             with a write method, or None. If a file-like
                             object, command  output will be written using its
                             write() method.
            * returnoutput - If True, output will be returned by the function.
                             Note that, when True, an exception will be raised
                             if output exceeds maxOutputSize. Defaults to True.
            * timeout      - A float value specifying the amount of time,
                             in seconds, to wait for output or completion. If
                             the timeout expires before output is read or the
                             process terminates, an exception is raised. A
                             value of 0 disables the timeout, and this function
                             may block indefinitely. Defaults to 0.
            * redirectErr  - If True, the error output will be redirected to the
                             output.
            * maxOutputSize - Maximum output size in bytes. If None is passed,
                              there is no limit.
            * stdinStr     - string that should be sent to stdin for the
                             subprocess/command.
         Returns: A tuple of the command's exit value and output if the redirectErr
                  is true and a tuple of the command's exit value, output and error
                  if the redirectErr is false
         Raises:
            * RunCommandError - All errors are caught and raised as this
                                exception. May happen if Popen fails (e.g.
                                command or process image does not exist or
                                cannot be executed), if there is an I/O error
                                reading output, or if the output exceeds 1024
                                characters when returnoutput is True.
   """

   log.info("runcommand called with: args = %r, outfile = %r, "
            "returnoutput = %r, timeout = %r.",
            args, outfile, returnoutput, timeout)

   try:
      if redirectErr:
         stdErr = subprocess.STDOUT
      else:
         stdErr = subprocess.PIPE

      stdinread = None
      if stdinStr:
         stdinread, stdinwrite = os.pipe()

      if ONVISOR and HAVEVMKUSER:
         p = VisorPopen(args, redirectErr, stdin=stdinread)
      elif isString(args):
         p = subprocess.Popen(args, stdin=stdinread, stdout=subprocess.PIPE,
                              stderr=stdErr, shell=True)
      else:
         p = subprocess.Popen(args, stdin=stdinread, stdout=subprocess.PIPE,
                              stderr=stdErr)
      if stdinread:
         os.close(stdinread)
         stdinWriteFP = os.fdopen(stdinwrite, 'w')
         stdinWriteFP.write(stdinStr)
         stdinWriteFP.close()
   except Exception as e:
      raise RunCommandError("Error running command %r: %s" % (args, e)) from e

   return waitProcessToComplete(p, args, outfile, returnoutput,
                                timeout, maxOutputSize)

def waitProcessToComplete( p, args,
                           outfile=None,
                           returnoutput=True,
                           timeout=0.0,
                           maxOutputSize=MAXOUTPUTSIZE):
   pollobj = select.poll()
   pollobj.register(p.stdout, select.POLLIN|select.POLLPRI)

   for pipe in [p.stdout, p.stderr]:
      if pipe is None:
         continue

      # Set O_NONBLOCK on pipe, else we can hang on a read. See PR 110272.
      flags = fcntl.fcntl(pipe, fcntl.F_GETFL) | os.O_NONBLOCK
      fcntl.fcntl(pipe, fcntl.F_SETFL, flags)
      pollobj.register(pipe, select.POLLIN|select.POLLPRI)

   data = b''
   out = []
   if p.stderr:
      err = []
      fd_dict = {p.stdout.fileno(): (p.stdout, out),
                 p.stderr.fileno(): (p.stderr, err)}
   else:
      err = out
      fd_dict = {p.stdout.fileno(): (p.stdout, out)}

   while p.poll() == None:
      try:
         if timeout > 0:
            ready = pollobj.poll(timeout * 1000)
         else:
            # See the "baddaemon" test in test/test_runcommand.py. By using a
            # timeout here, it gives us a chance to re-evaluate the conditional
            # in the while loop. This will catch daemon processes which don't
            # close their stdin/stdout/stderr.
            ready = pollobj.poll(200)
      except select.error as e:
         (n, strerror) = e
         # If we're interrupted by a signal, that's ok. Otherwise, it's
         # something bad and we should bail.
         if n == errno.EINTR:
            continue
         msg = ("Error waiting for output from command '%s': %s" %
                (args, strerror))
         raise RunCommandError(msg)
      except Exception as e:
         # Bail out and terminate process if there're unhandled exceptions
         terminateprocess(p)
         msg = ("Exception '%s' occured waiting for output from command '%s', "
                "pid '%s'" % (type(e).__name__, args, p.pid))
         raise RunCommandError(msg)
      if timeout and not ready and p.poll() == None:
         # The select() call returned before we had a signal or output, and the
         # command hasn't finished yet.
         terminateprocess(p)
         msg = ("Timeout (%d seconds) expired waiting for output from "
                "command '%s', pid '%s'." % (timeout, args, p.pid))
         raise RunCommandError(msg)
      try:
         # We read up to 1024 bytes at a time, to avoid usurping lots of
         # memory. If there is more data in the pipe, the next select will just
         # return right away.
         for fd, event in ready:
            fdPipe = fd_dict[fd][0]
            data = fdPipe.read(1024)
            # Skip the data with zero data length to avoid potential memory leak
            # when adding the 0 byte data into output buffer
            if len(data) == 0:
               continue

            # If we've read data, send it to outfile and/or add it to the collected
            # data to return.
            if outfile:
               outfile.write(data)
            if returnoutput:
               output = fd_dict[fd][1]
               if maxOutputSize is not None and \
                  sum(len(s) for s in output) + len(data) > maxOutputSize:
                  msg = "Output from command '%s' is too large and exceeds %d bytes" % \
                        (args, maxOutputSize)
                  raise RunCommandError(msg)
               output.append(data)
      except IOError as e:
         if e.errno == errno.EAGAIN:
            # We've set O_NONBLOCK, so the read() may return EAGAIN.
            continue
         else:
            msg = "Error reading output from command '%s': %s" % (args, e)
            raise RunCommandError(msg)

   def readDataFromPipe(pipe, output):
      # Read in any more data in the pipe.
      try:
         data = pipe.read(1024)
      except IOError:
         data = b''
      while data:
         if len(data) > 0:
            if outfile:
               outfile.write(data)
            if returnoutput:
               if maxOutputSize is not None and \
                  sum(len(s) for s in output) + len(data) > maxOutputSize:
                  msg = "Output from command '%s' is too large and exceeds %d bytes" % \
                        (args, maxOutputSize)
                  raise RunCommandError(msg)
               output.append(data)
         try:
            data = pipe.read(1024)
         except IOError:
            data = b''

   readDataFromPipe(p.stdout, out)
   if p.stderr:
      readDataFromPipe(p.stderr, err)

   # Get the status of the child.
   rc = p.poll()

   # Note: subprocess.Popen's poll() method will return a negative value if the
   #       process exited due to a signal.
   if rc >= 0:
      output = b"".join(out)

      if p.stderr:
         errput = b"".join(err)
         return rc, output, errput
      else:
         return rc, output
   else:
      # If we exited due to a signal, raise an exception.
      msg = "Command '%s' terminated due to signal %d" % (args, -rc)
      raise RunCommandError(msg)
