pec2 0.3.2 → 0.4.0

data/exe/psshlib/psshutil.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2009-2012, Andrew McNabb
+# Copyright (c) 2003-2008, Brent N. Chun
+
+import fcntl
+import string
+import sys
+
+HOST_FORMAT = 'Host format is [user@]host[:port] [user]'
+
+
+def read_host_files(paths, default_user=None, default_port=None):
+    """Reads the given host files.
+
+    Returns a list of (host, port, user) triples.
+    """
+    hosts = []
+    if paths:
+        for path in paths:
+            hosts.extend(read_host_file(path, default_user=default_user))
+    return hosts
+
+
+def read_host_file(path, default_user=None, default_port=None):
+    """Reads the given host file.
+
+    Lines are of the form: host[:port] [login].
+    Returns a list of (host, port, user) triples.
+    """
+    lines = []
+    f = open(path)
+    for line in f:
+        lines.append(line.strip())
+    f.close()
+
+    hosts = []
+    for line in lines:
+        # Skip blank lines or lines starting with #
+        line = line.strip()
+        if not line or line.startswith('#'):
+            continue
+        host, port, user = parse_host_entry(line, default_user, default_port)
+        if host:
+            hosts.append((host, port, user))
+    return hosts
+
+
+# TODO: deprecate the second host field and standardize on the
+# [user@]host[:port] format.
+def parse_host_entry(line, default_user, default_port):
+    """Parses a single host entry.
+
+    This may take either the of the form [user@]host[:port] or
+    host[:port][ user].
+
+    Returns a (host, port, user) triple.
+    """
+    fields = line.split()
+    if len(fields) > 2:
+        sys.stderr.write('Bad line: "%s". Format should be'
+                ' [user@]host[:port] [user]\n' % line)
+        return None, None, None
+    host_field = fields[0]
+    host, port, user = parse_host(host_field, default_port=default_port)
+    if len(fields) == 2:
+        if user is None:
+            user = fields[1]
+        else:
+            sys.stderr.write('User specified twice in line: "%s"\n' % line)
+            return None, None, None
+    if user is None:
+        user = default_user
+    return host, port, user
+
+
+def parse_host_string(host_string, default_user=None, default_port=None):
+    """Parses a whitespace-delimited string of "[user@]host[:port]" entries.
+
+    Returns a list of (host, port, user) triples.
+    """
+    hosts = []
+    entries = host_string.split()
+    for entry in entries:
+        hosts.append(parse_host(entry, default_user, default_port))
+    return hosts
+
+
+def parse_host(host, default_user=None, default_port=None):
+    """Parses host entries of the form "[user@]host[:port]".
+
+    Returns a (host, port, user) triple.
+    """
+    # TODO: when we stop supporting Python 2.4, switch to using str.partition.
+    user = default_user
+    port = default_port
+    if '@' in host:
+        user, host = host.split('@', 1)
+    if ':' in host:
+        host, port = host.rsplit(':', 1)
+    return (host, port, user)
+
+
+def set_cloexec(filelike):
+    """Sets the underlying filedescriptor to automatically close on exec.
+
+    If set_cloexec is called for all open files, then subprocess.Popen does
+    not require the close_fds option.
+    """
+    fcntl.fcntl(filelike.fileno(), fcntl.FD_CLOEXEC, 1)

data/exe/psshlib/task.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2009-2012, Andrew McNabb
+
+from errno import EINTR
+from subprocess import Popen, PIPE
+import os
+import signal
+import sys
+import time
+import traceback
+
+from psshlib import askpass_client
+from psshlib import color
+
+BUFFER_SIZE = 1 << 16
+
+try:
+    bytes
+except NameError:
+    bytes = str
+
+
+class Task(object):
+    """Starts a process and manages its input and output.
+
+    Upon completion, the `exitstatus` attribute is set to the exit status
+    of the process.
+    """
+    def __init__(self, host, port, user, cmd, opts, stdin=None):
+        self.exitstatus = None
+
+        self.host = host
+        self.pretty_host = host
+        self.port = port
+        self.cmd = cmd
+
+        if user != opts.user:
+            self.pretty_host = '@'.join((user, self.pretty_host))
+        if port:
+            self.pretty_host = ':'.join((self.pretty_host, port))
+
+        self.proc = None
+        self.writer = None
+        self.timestamp = None
+        self.failures = []
+        self.killed = False
+        self.inputbuffer = stdin
+        self.byteswritten = 0
+        self.outputbuffer = bytes()
+        self.errorbuffer = bytes()
+
+        self.stdin = None
+        self.stdout = None
+        self.stderr = None
+        self.outfile = None
+        self.errfile = None
+
+        # Set options.
+        self.verbose = opts.verbose
+        try:
+            self.print_out = bool(opts.print_out)
+        except AttributeError:
+            self.print_out = False
+        try:
+            self.inline = bool(opts.inline)
+        except AttributeError:
+            self.inline = False
+        try:
+            self.inline_stdout = bool(opts.inline_stdout)
+        except AttributeError:
+            self.inline_stdout = False
+
+    def start(self, nodenum, iomap, writer, askpass_socket=None):
+        """Starts the process and registers files with the IOMap."""
+        self.writer = writer
+
+        if writer:
+            self.outfile, self.errfile = writer.open_files(self.pretty_host)
+
+        # Set up the environment.
+        environ = dict(os.environ)
+        environ['PSSH_NODENUM'] = str(nodenum)
+        environ['PSSH_HOST'] = self.host
+        # Disable the GNOME pop-up password dialog and allow ssh to use
+        # askpass.py to get a provided password.  If the module file is
+        # askpass.pyc, we replace the extension.
+        environ['SSH_ASKPASS'] = askpass_client.executable_path()
+        if askpass_socket:
+            environ['PSSH_ASKPASS_SOCKET'] = askpass_socket
+        if self.verbose:
+            environ['PSSH_ASKPASS_VERBOSE'] = '1'
+        # Work around a mis-feature in ssh where it won't call SSH_ASKPASS
+        # if DISPLAY is unset.
+        if 'DISPLAY' not in environ:
+            environ['DISPLAY'] = 'pssh-gibberish'
+
+        # Create the subprocess.  Since we carefully call set_cloexec() on
+        # all open files, we specify close_fds=False.
+        self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE,
+                close_fds=False, preexec_fn=os.setsid, env=environ)
+        self.timestamp = time.time()
+        if self.inputbuffer:
+            self.stdin = self.proc.stdin
+            iomap.register_write(self.stdin.fileno(), self.handle_stdin)
+        else:
+            self.proc.stdin.close()
+        self.stdout = self.proc.stdout
+        iomap.register_read(self.stdout.fileno(), self.handle_stdout)
+        self.stderr = self.proc.stderr
+        iomap.register_read(self.stderr.fileno(), self.handle_stderr)
+
+    def _kill(self):
+        """Signals the process to terminate."""
+        if self.proc:
+            try:
+                os.kill(-self.proc.pid, signal.SIGKILL)
+            except OSError:
+                # If the kill fails, then just assume the process is dead.
+                pass
+            self.killed = True
+
+    def timedout(self):
+        """Kills the process and registers a timeout error."""
+        if not self.killed:
+            self._kill()
+            self.failures.append('Timed out')
+
+    def interrupted(self):
+        """Kills the process and registers an keyboard interrupt error."""
+        if not self.killed:
+            self._kill()
+            self.failures.append('Interrupted')
+
+    def cancel(self):
+        """Stops a task that has not started."""
+        self.failures.append('Cancelled')
+
+    def elapsed(self):
+        """Finds the time in seconds since the process was started."""
+        return time.time() - self.timestamp
+
+    def running(self):
+        """Finds if the process has terminated and saves the return code."""
+        if self.stdin or self.stdout or self.stderr:
+            return True
+        if self.proc:
+            self.exitstatus = self.proc.poll()
+            if self.exitstatus is None:
+                if self.killed:
+                    # Set the exitstatus to what it would be if we waited.
+                    self.exitstatus = -signal.SIGKILL
+                    return False
+                else:
+                    return True
+            else:
+                if self.exitstatus < 0:
+                    message = 'Killed by signal %s' % (-self.exitstatus)
+                    self.failures.append(message)
+                elif self.exitstatus > 0:
+                    message = 'Exited with error code %s' % self.exitstatus
+                    self.failures.append(message)
+                self.proc = None
+                return False
+
+    def handle_stdin(self, fd, iomap):
+        """Called when the process's standard input is ready for writing."""
+        try:
+            start = self.byteswritten
+            if start < len(self.inputbuffer):
+                chunk = self.inputbuffer[start:start+BUFFER_SIZE]
+                self.byteswritten = start + os.write(fd, chunk)
+            else:
+                self.close_stdin(iomap)
+        except (OSError, IOError):
+            _, e, _ = sys.exc_info()
+            if e.errno != EINTR:
+                self.close_stdin(iomap)
+                self.log_exception(e)
+
+    def close_stdin(self, iomap):
+        if self.stdin:
+            iomap.unregister(self.stdin.fileno())
+            self.stdin.close()
+            self.stdin = None
+
+    def handle_stdout(self, fd, iomap):
+        """Called when the process's standard output is ready for reading."""
+        try:
+            buf = os.read(fd, BUFFER_SIZE)
+            if buf:
+                if self.inline or self.inline_stdout:
+                    self.outputbuffer += buf
+                if self.outfile:
+                    self.writer.write(self.outfile, buf)
+                if self.print_out:
+                    sys.stdout.write('%s: %s' % (self.host, buf))
+                    if buf[-1] != '\n':
+                        sys.stdout.write('\n')
+            else:
+                self.close_stdout(iomap)
+        except (OSError, IOError):
+            _, e, _ = sys.exc_info()
+            if e.errno != EINTR:
+                self.close_stdout(iomap)
+                self.log_exception(e)
+
+    def close_stdout(self, iomap):
+        if self.stdout:
+            iomap.unregister(self.stdout.fileno())
+            self.stdout.close()
+            self.stdout = None
+        if self.outfile:
+            self.writer.close(self.outfile)
+            self.outfile = None
+
+    def handle_stderr(self, fd, iomap):
+        """Called when the process's standard error is ready for reading."""
+        try:
+            buf = os.read(fd, BUFFER_SIZE)
+            if buf:
+                if self.inline:
+                    self.errorbuffer += buf
+                if self.errfile:
+                    self.writer.write(self.errfile, buf)
+            else:
+                self.close_stderr(iomap)
+        except (OSError, IOError):
+            _, e, _ = sys.exc_info()
+            if e.errno != EINTR:
+                self.close_stderr(iomap)
+                self.log_exception(e)
+
+    def close_stderr(self, iomap):
+        if self.stderr:
+            iomap.unregister(self.stderr.fileno())
+            self.stderr.close()
+            self.stderr = None
+        if self.errfile:
+            self.writer.close(self.errfile)
+            self.errfile = None
+
+    def log_exception(self, e):
+        """Saves a record of the most recent exception for error reporting."""
+        if self.verbose:
+            exc_type, exc_value, exc_traceback = sys.exc_info()
+            exc = ("Exception: %s, %s, %s" %
+                    (exc_type, exc_value, traceback.format_tb(exc_traceback)))
+        else:
+            exc = str(e)
+        self.failures.append(exc)
+
+    def report(self, n):
+        """Pretty prints a status report after the Task completes."""
+        error = ', '.join(self.failures)
+        tstamp = time.asctime().split()[3] # Current time
+        if color.has_colors(sys.stdout):
+            progress = color.c("[%s]" % color.B(n))
+            success = color.g("[%s]" % color.B("SUCCESS"))
+            failure = color.r("[%s]" % color.B("FAILURE"))
+            stderr = color.r("Stderr: ")
+            error = color.r(color.B(error))
+        else:
+            progress = "[%s]" % n
+            success = "[SUCCESS]"
+            failure = "[FAILURE]"
+            stderr = "Stderr: "
+        host = self.pretty_host
+        if self.failures:
+            print(' '.join((progress, tstamp, failure, host, error)))
+        else:
+            print(' '.join((progress, tstamp, success, host)))
+        # NOTE: The extra flushes are to ensure that the data is output in
+        # the correct order with the C implementation of io.
+        if self.outputbuffer:
+            sys.stdout.flush()
+            try:
+                sys.stdout.buffer.write(self.outputbuffer)
+                sys.stdout.flush()
+            except AttributeError:
+                sys.stdout.write(self.outputbuffer)
+        if self.errorbuffer:
+            sys.stdout.write(stderr)
+            # Flush the TextIOWrapper before writing to the binary buffer.
+            sys.stdout.flush()
+            try:
+                sys.stdout.buffer.write(self.errorbuffer)
+            except AttributeError:
+                sys.stdout.write(self.errorbuffer)
+

data/exe/psshlib/version.py

@@ -0,0 +1 @@
+VERSION = '2.3.1'

data/lib/pec2/cli.rb

@@ -1,5 +1,7 @@
 require "thor"
 require "tempfile"
+require "logger"
+require 'shellwords'
 
 module Pec2
   class CLI < Thor
@@ -11,7 +13,8 @@ module Pec2
       super(args, options, config)
       @global_options = config[:shell].base.options
       @core = Core.new
-      @pssh_path = `which pssh`.strip
+      @pssh_path = File.expand_path('../../../exe/bin/pssh', __FILE__)
+      @logger = Logger.new(STDOUT)
     end
 
     desc 'search_tag', 'search tag'
@@ -29,7 +32,8 @@ module Pec2
         end
 
         if addresses.empty?
-          raise "no host."
+          @logger.error(%Q{no host tag #{options[:tag]}.})
+          raise
         end
 
         File.write(f.path, addresses.join("\n"))
@@ -52,9 +56,9 @@ module Pec2
         end
 
         if options[:sudo_password]
-          cmd = %Q{(echo #{options[:sudo_password]}) | #{cmd} -I '#{options[:command]}'}
+          cmd = %Q{(echo #{options[:sudo_password]}) | #{cmd} -I #{Shellwords.escape(options[:command])}}
         else
-          cmd = %Q{#{cmd} -i '#{options[:command]}'}
+          cmd = %Q{#{cmd} -i #{Shellwords.escape(options[:command])}}
         end
         system(cmd)
       end

data/lib/pec2/core.rb

@@ -6,7 +6,7 @@ require 'json'
 require 'timeout'
 
 class Pec2Mash < ::Hashie::Mash
-  disable_warnings
+  disable_warnings if respond_to?(:disable_warnings)
 end
 
 module Pec2