extended usage of Executor class, added client-side timeout options

1 files changed, 155 insertions, 1 deletions

diff --git a/src/lib/Bcfg2/Utils.py b/src/lib/Bcfg2/Utils.py
index 247e4f16b..3b1559528 100644
--- a/src/lib/Bcfg2/Utils.py
+++ b/src/lib/Bcfg2/Utils.py
@@ -3,9 +3,25 @@ used by both client and server.  Stuff that doesn't fit anywhere
 else. """
 
 import fcntl
+import logging
+import threading
+import subprocess
 from Bcfg2.Compat import any  # pylint: disable=W0622
 
 
+class ClassName(object):
+    """ This very simple descriptor class exists only to get the name
+    of the owner class.  This is used because, for historical reasons,
+    we expect every server plugin and every client tool to have a
+    ``name`` attribute that is in almost all cases the same as the
+    ``__class__.__name__`` attribute of the plugin object.  This makes
+    that more dynamic so that each plugin and tool isn't repeating its own
+    name."""
+
+    def __get__(self, inst, owner):
+        return owner.__name__
+
+
 class PackedDigitRange(object):
     """ Representation of a set of integer ranges. A range is
     described by a comma-delimited string of integers and ranges,
@@ -69,9 +85,147 @@ class PackedDigitRange(object):
 
 
 def locked(fd):
-    """ Acquire a lock on a file """
+    """ Acquire a lock on a file.
+
+    :param fd: The file descriptor to lock
+    :type fd: int
+    :returns: bool - True if the file is already locked, False
+              otherwise """
     try:
         fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
     except IOError:
         return True
     return False
+
+
+class ExecutorResult(object):
+    """ Returned as the result of a call to
+    :func:`Bcfg2.Utils.Executor.run`. The result can be accessed via
+    the instance variables, documented below, as a boolean (which is
+    equivalent to :attr:`Bcfg2.Utils.ExecutorResult.success`), or as a
+    tuple, which, for backwards compatibility, is equivalent to
+    ``(result.retval, result.stdout.splitlines())``."""
+
+    def __init__(self, stdout, stderr, retval):
+        #: The output of the command
+        self.stdout = stdout
+
+        #: The error produced by the command
+        self.stderr = stderr
+
+        #: The return value of the command.
+        self.retval = retval
+
+        #: Whether or not the command was successful.  If the
+        #: ExecutorResult is used as a boolean, ``success`` is
+        #: returned.
+        self.success = retval == 0
+
+        #: A friendly error message
+        self.error = None
+        if self.retval:
+            if self.stderr:
+                self.error = "%s (rv: %s)" % (self.stderr, self.retval)
+            elif self.stdout:
+                self.error = "%s (rv: %s)" % (self.stdout, self.retval)
+            else:
+                self.error = "No output or error; return value %s" % \
+                    self.retval
+
+    def __repr__(self):
+        if self.error:
+            return "Errored command result: %s" % self.error
+        elif self.stdout:
+            return "Successful command result: %s" % self.stdout
+        else:
+            return "Successful command result: No output"
+
+    def __getitem__(self, idx):
+        """ This provides compatibility with the old Executor, which
+        returned a tuple of (return value, stdout split by lines). """
+        return (self.retval, self.stdout.splitlines())[idx]
+
+    def __nonzero__(self):
+        return self.__bool__()
+
+    def __bool__(self):
+        return self.success
+
+
+class Executor(object):
+    """ A convenient way to run external commands with
+    :class:`subprocess.Popen` """
+
+    def __init__(self, timeout=None):
+        """
+        :param timeout: Set a default timeout for all commands run by
+                        this Executor object
+        :type timeout: float
+        """
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.timeout = timeout
+
+    def _timeout_callback(self, proc):
+        """ Get a callback (suitable for passing to
+        :class:`threading.Timer`) that kills the given process.
+
+        :param proc: The process to kill upon timeout.
+        :type proc: subprocess.Popen
+        :returns: function """
+        def _timeout():
+            """ Callback that kills ``proc`` """
+            if proc.poll() == None:
+                try:
+                    proc.kill()
+                    self.logger.warning("Process exceeeded timeout, killing")
+                except OSError:
+                    pass
+
+        return _timeout
+
+    def run(self, command, inputdata=None, shell=False, timeout=None):
+        """ Run a command, given as a list, optionally giving it the
+        specified input data.
+
+        :param command: The command to run, as a list (preferred) or
+                        as a string.  See :class:`subprocess.Popen` for
+                        details.
+        :type command: list or string
+        :param inputdata: Data to pass to the command on stdin
+        :type inputdata: string
+        :param shell: Run the given command in a shell (not recommended)
+        :type shell: bool
+        :param timeout: Kill the command if it runs longer than this
+                        many seconds.  Set to 0 or -1 to explicitly
+                        override a default timeout.
+        :type timeout: float
+        :returns: :class:`Bcfg2.Utils.ExecutorResult`
+        """
+        if isinstance(command, str):
+            cmdstr = command
+        else:
+            cmdstr = " ".join(command)
+        self.logger.debug("Running: %s" % cmdstr)
+        proc = subprocess.Popen(command, shell=shell, bufsize=16384,
+                                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                                stderr=subprocess.PIPE, close_fds=True)
+        if timeout is None:
+            timeout = self.timeout
+        if timeout is not None:
+            timer = threading.Timer(float(timeout),
+                                    self._timeout_callback(proc),
+                                    [proc])
+            timer.start()
+        try:
+            if inputdata:
+                for line in inputdata.splitlines():
+                    self.logger.debug('> %s' % line)
+            (stdout, stderr) = proc.communicate(input=inputdata)
+            for line in stdout.splitlines():  # pylint: disable=E1103
+                self.logger.debug('< %s' % line)
+            for line in stderr.splitlines():  # pylint: disable=E1103
+                self.logger.info(line)
+            return ExecutorResult(stdout, stderr, proc.wait())
+        finally:
+            if timeout is not None:
+                timer.cancel()