capissh 0.0.1

data/lib/capissh/connection_manager.rb

@@ -0,0 +1,250 @@
+require 'enumerator'
+require 'net/ssh/gateway'
+require 'capissh/ssh'
+require 'capissh/errors'
+require 'capissh/server_definition'
+require 'capissh/logger'
+
+module Capissh
+  class ConnectionManager
+    class DefaultConnectionFactory
+      def initialize(options)
+        @options = options
+      end
+
+      def connect_to(server)
+        SSH.connect(server, @options)
+      end
+    end
+
+    class GatewayConnectionFactory
+      def initialize(gateway, options)
+        @options = options
+        Thread.abort_on_exception = true
+        @gateways = {}
+        @default_gateway = nil
+        if gateway.is_a?(Hash)
+          @options[:logger].debug "Creating multiple gateways using #{gateway.inspect}" if @options[:logger]
+          gateway.each do |gw, hosts|
+            gateway_connection = add_gateway(gw)
+            Array(hosts).each do |host|
+              # Why is the default only set if there's at least one host?
+              # It seems odd enough to be intentional, since it could easily be set outside of the loop.
+              @default_gateway ||= gateway_connection
+              @gateways[host] = gateway_connection
+            end
+          end
+        else
+          @options[:logger].debug "Creating gateway using #{Array(gateway).join(', ')}" if @options[:logger]
+          @default_gateway = add_gateway(gateway)
+        end
+      end
+
+      def add_gateway(gateway)
+        gateways = ServerDefinition.wrap_list(gateway)
+
+        tunnel = SSH.gateway(gateways.shift, @options)
+
+        gateways.inject(tunnel) do |tunnel, destination|
+          @options[:logger].debug "Creating tunnel to #{destination}" if @options[:logger]
+          local = local_host(destination, tunnel)
+          SSH.gateway(local, @options)
+        end
+      end
+
+      def connect_to(server)
+        @options[:logger].debug "establishing connection to `#{server}' via gateway" if @options[:logger]
+        local = local_host(server, gateway_for(server))
+        session = SSH.connect(local, @options)
+        session.xserver = server
+        session
+      end
+
+      def local_host(destination, tunnel)
+        ServerDefinition.new("127.0.0.1", :user => destination.user, :port => tunnel.open(destination.host, destination.connect_to_port))
+      end
+
+      def gateway_for(server)
+        @gateways[server.host] || @default_gateway
+      end
+    end
+
+    # Instantiates a new ConnectionManager object.
+    # +options+ must be a hash containing any of the following keys:
+    #
+    # * +gateway+: (optional), ssh gateway
+    # * +logger+: (optional), a Capissh::Logger instance
+    # * +ssh_options+: (optional), options for Net::SSH
+    # * +verbose+: (optional), verbosity level for ssh
+    # * +user+: (optional), user for all servers
+    # * +port+: (optional), port for all servers
+    # * +password+: (optional), password for all servers
+    #
+    # * many more that I haven't written yet
+    #
+    def initialize(options={})
+      @options = options.dup
+      @gateway = @options.delete(:gateway)
+      @logger  = @options[:logger] || Capissh::Logger.default
+      Thread.current[:sessions] = {}
+      Thread.current[:failed_sessions] = []
+    end
+
+    attr_reader :logger
+
+    # A hash of the SSH sessions that are currently open and available.
+    # Because sessions are constructed lazily, this will only contain
+    # connections to those servers that have been the targets of one or more
+    # executed tasks. Stored on a per-thread basis to improve thread-safety.
+    def sessions
+      Thread.current[:sessions] ||= {}
+    end
+
+    # Indicate that the given server could not be connected to.
+    def failed!(server)
+      Thread.current[:failed_sessions] << server
+    end
+
+    # Query whether previous connection attempts to the given server have
+    # failed.
+    def has_failed?(server)
+      Thread.current[:failed_sessions].include?(server)
+    end
+
+    # Used to force connections to be made to the current task's servers.
+    # Connections are normally made lazily in Capissh--you can use this
+    # to force them open before performing some operation that might be
+    # time-sensitive.
+    def connect!(servers, options={})
+      execute_on_servers(servers, options) { }
+    end
+
+    # Returns the object responsible for establishing new SSH connections.
+    # The factory will respond to #connect_to, which can be used to
+    # establish connections to servers defined via ServerDefinition objects.
+    def connection_factory
+      @connection_factory ||= begin
+        if @gateway
+          logger.debug "establishing connection to gateway `#{@gateway.inspect}'"
+          GatewayConnectionFactory.new(@gateway, @options)
+        else
+          DefaultConnectionFactory.new(@options)
+        end
+      end
+    end
+
+    # Ensures that there are active sessions for each server in the list.
+    def establish_connections_to(servers)
+      # force the connection factory to be instantiated synchronously,
+      # otherwise we wind up with multiple gateway instances, because
+      # each connection is done in parallel.
+      connection_factory
+
+      failed_servers = []
+      servers_array = Array(servers)
+
+      threads = servers_array.map { |server| establish_connection_to(server, failed_servers) }
+      threads.each { |t| t.join }
+
+      if failed_servers.any?
+        messages = failed_servers.map { |h| "#{h[:server]} (#{h[:error].class}: #{h[:error].message})" }
+        error = ConnectionError.new("connection failed for: #{messages.join(', ')}")
+        error.hosts = failed_servers.map { |h| h[:server] }.each {|server| failed!(server) }
+        raise error
+      end
+
+      servers_array.map {|server| sessions[server] }
+    end
+
+    # Destroys sessions for each server in the list.
+    def teardown_connections_to(servers)
+      servers.each do |server|
+        begin
+          sessions.delete(server).close
+        rescue IOError
+          # the TCP connection is already dead
+        end
+      end
+    end
+
+    # Determines the set of servers and establishes connections to them,
+    # and then yields that list of servers.
+    #
+    # All options will be used to find servers. (see find_servers)
+    #
+    # The additional options below will also be used as follows:
+    #
+    # * +on_no_matching_servers+: (optional), set to :continue to return
+    #   instead of raising when no servers are found
+    # * +max_hosts+: (optional), integer to batch commands in chunks of hosts
+    # * +continue_on_error+: (optionsal), continue on connection errors
+    #
+    def execute_on_servers(servers, options={}, &block)
+      raise ArgumentError, "expected a block" unless block_given?
+
+      connect_to_servers = ServerDefinition.wrap_list(*servers)
+
+      if options[:continue_on_error]
+        connect_to_servers.delete_if { |s| has_failed?(s) }
+      end
+
+      if connect_to_servers.empty?
+        raise Capissh::NoMatchingServersError, "no servers found to match #{options.inspect}" unless options[:continue_on_no_matching_servers]
+        return
+      end
+
+      logger.trace "servers: #{connect_to_servers.map { |s| s.host }.inspect}"
+
+      max_hosts = (options[:max_hosts] || connect_to_servers.size).to_i
+      is_subset = max_hosts < connect_to_servers.size
+
+      if max_hosts <= 0
+        raise Capissh::NoMatchingServersError, "max_hosts is invalid for #{options.inspect}" unless options[:continue_on_no_matching_servers]
+        return
+      end
+
+      # establish connections to those servers in groups of max_hosts, as necessary
+      connect_to_servers.each_slice(max_hosts) do |servers_slice|
+        begin
+          establish_connections_to(servers_slice)
+        rescue ConnectionError => error
+          raise error unless options[:continue_on_error]
+          error.hosts.each do |h|
+            servers_slice.delete(h)
+            failed!(h)
+          end
+        end
+
+        begin
+          yield servers_slice.map { |server| sessions[server] }
+        rescue RemoteError => error
+          raise error unless options[:continue_on_error]
+          error.hosts.each { |h| failed!(h) }
+        end
+
+        # if dealing with a subset (e.g., :max_hosts is less than the
+        # number of servers available) teardown the subset of connections
+        # that were just made, so that we can make room for the next subset.
+        teardown_connections_to(servers_slice) if is_subset
+      end
+    end
+
+    private
+
+    # We establish the connection by creating a thread in a new method--this
+    # prevents problems with the thread's scope seeing the wrong 'server'
+    # variable if the thread just happens to take too long to start up.
+    def establish_connection_to(server, failures=nil)
+      current_thread = Thread.current
+      Thread.new { safely_establish_connection_to(server, current_thread, failures) }
+    end
+
+    def safely_establish_connection_to(server, thread, failures=nil)
+      thread[:sessions] ||= {} # can this move up to the current_thread part above?
+      thread[:sessions][server] ||= connection_factory.connect_to(server)
+    rescue Exception => err
+      raise unless failures
+      failures << { :server => server, :error => err }
+    end
+  end
+end

data/lib/capissh/errors.rb

@@ -0,0 +1,19 @@
+module Capissh
+
+  Error = Class.new(RuntimeError)
+
+  CaptureError            = Class.new(Capissh::Error)
+  NoSuchTaskError         = Class.new(Capissh::Error)
+  NoMatchingServersError  = Class.new(Capissh::Error)
+
+  class RemoteError < Error
+    attr_accessor :hosts
+  end
+
+  ConnectionError     = Class.new(Capissh::RemoteError)
+  TransferError       = Class.new(Capissh::RemoteError)
+  CommandError        = Class.new(Capissh::RemoteError)
+
+  LocalArgumentError  = Class.new(Capissh::Error)
+
+end

data/lib/capissh/file_transfers.rb

@@ -0,0 +1,54 @@
+require 'capissh/transfer'
+
+module Capissh
+  class FileTransfers
+    attr_reader :configuration, :connection_manager, :logger
+
+    def initialize(configuration, connection_manager, options={})
+      @configuration = configuration
+      @connection_manager = connection_manager
+      @logger  = options[:logger]
+    end
+
+    # Store the given data at the given location on all servers targetted
+    # by the current task. If <tt>:mode</tt> is specified it is used to
+    # set the mode on the file.
+    def put(servers, data, path, options={})
+      upload(servers, StringIO.new(data), path, options)
+    end
+
+    # Get file remote_path from FIRST server targeted by
+    # the current task and transfer it to local machine as path.
+    #
+    # Pass only one server, or the first of the set of servers will be used.
+    #
+    # get server, "#{deploy_to}/current/log/production.log", "log/production.log.web"
+    def get(servers, remote_path, path, options={}, &block)
+      download(Array(servers).slice(0,1), remote_path, path, options, &block)
+    end
+
+    def upload(servers, from, to, options={}, &block)
+      opts = options.dup
+      mode = opts.delete(:mode)
+      transfer(servers, :up, from, to, opts, &block)
+      if mode
+        mode = mode.is_a?(Numeric) ? mode.to_s(8) : mode.to_s
+        configuration.run servers, "chmod #{mode} #{to}", opts
+      end
+    end
+
+    def download(servers, from, to, options={}, &block)
+      transfer(servers, :down, from, to, options, &block)
+    end
+
+    def transfer(servers, direction, from, to, options={}, &block)
+      if configuration.dry_run
+        return logger.debug "transfering: #{[direction, from, to] * ', '}"
+      end
+      connection_manager.execute_on_servers(servers, options) do |sessions|
+        Transfer.process(direction, from, to, sessions, options.merge(:logger => logger), &block)
+      end
+    end
+
+  end
+end

data/lib/capissh/invocation.rb

@@ -0,0 +1,278 @@
+require 'capissh/command'
+
+module Capissh
+  class Invocation
+    attr_reader :configuration, :logger, :connection_manager
+
+    def initialize(configuration, connection_manager, options={})
+      @configuration = configuration
+      @connection_manager = connection_manager
+      @logger = options[:logger]
+    end
+
+    # Executes different commands in parallel. This is useful for commands
+    # that need to be different on different hosts, but which could be
+    # otherwise run in parallel.
+    #
+    # The +options+ parameter is currently unused.
+    #
+    # Example:
+    #
+    #   parallel do |session|
+    #     session.when "in?(:app)", "/path/to/restart/mongrel"
+    #     session.when "in?(:web)", "/path/to/restart/apache"
+    #     session.when "in?(:db)", "/path/to/restart/mysql"
+    #   end
+    #
+    # Each command may have its own callback block, for capturing and
+    # responding to output, with semantics identical to #run:
+    #
+    #   session.when "in?(:app)", "/path/to/restart/mongrel" do |ch, stream, data|
+    #     # ch is the SSH channel for this command, used to send data
+    #     #    back to the command (e.g. ch.send_data("password\n"))
+    #     # stream is either :out or :err, for which stream the data arrived on
+    #     # data is a string containing data sent from the remote command
+    #   end
+    #
+    # Also, you can specify a fallback command, to use when none of the
+    # conditions match a server:
+    #
+    #   session.else "/execute/something/else"
+    #
+    # The string specified as the first argument to +when+ may be any valid
+    # Ruby code. It has access to the following variables and methods:
+    #
+    # * +server+ is the ServerDefinition object for the server. This can be
+    #   used to get the host-name, etc.
+    # * +configuration+ is the current Capissh::Configuration object, which
+    #   you can use to get the value of variables, etc.
+    #
+    # For example:
+    #
+    #   session.when "server.host =~ /app/", "/some/command"
+    #   session.when "server.host == configuration[:some_var]", "/another/command"
+    #   session.when "server.role?(:web) || server.role?(:app)", "/more/commands"
+    #
+    # See #run for a description of the valid +options+.
+    def parallel(servers, options={})
+      raise ArgumentError, "parallel() requires a block" unless block_given?
+      tree = Command::Tree.new(configuration) { |t| yield t }
+      run_tree(servers, tree, options)
+    end
+
+    # Invokes the given command. If a +via+ key is given, it will be used
+    # to determine what method to use to invoke the command. It defaults
+    # to :run, but may be :sudo, or any other method that conforms to the
+    # same interface as run and sudo.
+    def invoke_command(servers, cmd, options={}, &block)
+      options = options.dup
+      via = options.delete(:via) || :run
+      send(via, servers, cmd, options, &block)
+    end
+
+    # Execute the given command on all servers that are the target of the
+    # current task. If a block is given, it is invoked for all output
+    # generated by the command, and should accept three parameters: the SSH
+    # channel (which may be used to send data back to the remote process),
+    # the stream identifier (<tt>:err</tt> for stderr, and <tt>:out</tt> for
+    # stdout), and the data that was received.
+    #
+    # The +options+ hash may include any of the following keys:
+    #
+    # * :on_no_matching_servers - if :continue, will continue to execute tasks if
+    #   no matching servers are found for the host criteria. The default is to raise
+    #   a NoMatchingServersError exception.
+    # * :max_hosts - specifies the maximum number of hosts that should be selected
+    #   at a time. If this value is less than the number of hosts that are selected
+    #   to run, then the hosts will be run in groups of max_hosts. The default is nil,
+    #   which indicates that there is no maximum host limit. Please note this does not
+    #   limit the number of SSH channels that can be open, only the number of hosts upon
+    #   which this will be called.
+    # * :shell - says which shell should be used to invoke commands. This
+    #   defaults to "sh". Setting this to false causes Capissh to invoke
+    #   the commands directly, without wrapping them in a shell invocation.
+    # * :data - if not nil (the default), this should be a string that will
+    #   be passed to the command's stdin stream.
+    # * :pty - if true, a pseudo-tty will be allocated for each command. The
+    #   default is false. Note that there are benefits and drawbacks both ways.
+    #   Empirically, it appears that if a pty is allocated, the SSH server daemon
+    #   will _not_ read user shell start-up scripts (e.g. bashrc, etc.). However,
+    #   if a pty is _not_ allocated, some commands will refuse to run in
+    #   interactive mode and will not prompt for (e.g.) passwords.
+    # * :env - a hash of environment variable mappings that should be made
+    #   available to the command. The keys should be environment variable names,
+    #   and the values should be their corresponding values. The default is
+    #   empty, but may be modified by changing the +default_environment+
+    #   Capissh variable.
+    # * :eof - if true, the standard input stream will be closed after sending
+    #   any data specified in the :data option. If false, the input stream is
+    #   left open. The default is to close the input stream only if no block is
+    #   passed.
+    #
+    # Note that if you set these keys in the +default_run_options+ Capissh
+    # variable, they will apply for all invocations of #run, #invoke_command,
+    # and #parallel.
+    def run(servers, cmd, options={}, &block)
+      if options[:eof].nil? && !cmd.include?(sudo_command)
+        options = options.merge(:eof => !block_given?)
+      end
+      block ||= Command.default_io_proc
+      tree = Command::Tree.twig(configuration, cmd, &block)
+      run_tree(servers, tree, options)
+    end
+
+    # Executes a Capissh::Command::Tree object. This is not for direct
+    # use, but should instead be called indirectly, via #run or #parallel,
+    # or #invoke_command.
+    def run_tree(servers, tree, options={})
+      if tree.branches.empty? && tree.fallback
+        logger.debug "executing #{tree.fallback}" unless options[:silent]
+      elsif tree.branches.any?
+        logger.debug "executing multiple commands in parallel"
+        tree.each do |branch|
+          logger.trace "-> #{branch}"
+        end
+      else
+        raise ArgumentError, "attempt to execute without specifying a command"
+      end
+
+      return if configuration.dry_run || (configuration.debug && continue_execution(tree) == false)
+
+      options = add_default_command_options(options)
+
+      tree.each do |branch|
+        if branch.command.include?(sudo_command)
+          branch.callback = sudo_behavior_callback(branch.callback)
+        end
+      end
+
+      connection_manager.execute_on_servers(servers, options) do |sessions|
+        Command.process(tree, sessions, options.merge(:logger => logger))
+      end
+    end
+
+    # Invoked like #run, but executing the command via sudo.
+    # This assumes that the sudo password (if required) is the same as the
+    # password for logging in to the server.
+    #
+    #   sudo "mkdir /path/to/dir"
+    #
+    # Also, this method understands a <tt>:sudo</tt> configuration variable,
+    # which (if specified) will be used as the full path to the sudo
+    # executable on the remote machine:
+    #
+    #   Capissh.new(sudo: "/opt/local/bin/sudo")
+    #
+    # If you know what you're doing, you can also set <tt>:sudo_prompt</tt>,
+    # which tells capissh which prompt sudo should use when asking for
+    # a password. (This is so that capissh knows what prompt to look for
+    # in the output.) If you set :sudo_prompt to an empty string, Capissh
+    # will not send a preferred prompt.
+    def sudo(servers, command, options={}, &block)
+      run(servers, "#{sudo_command(options)} #{command}", options, &block)
+    end
+
+    # Returns the command string used by capissh to invoke a comamnd via
+    # sudo.
+    #
+    #   run "#{sudo_command :as => 'bob'} mkdir /path/to/dir"
+    #
+    # Also, this method understands a <tt>:sudo</tt> configuration variable,
+    # which (if specified) will be used as the full path to the sudo
+    # executable on the remote machine:
+    #
+    #   Capissh.new(sudo: "/opt/local/bin/sudo")
+    #
+    # If you know what you're doing, you can also set <tt>:sudo_prompt</tt>,
+    # which tells capissh which prompt sudo should use when asking for
+    # a password. (This is so that capissh knows what prompt to look for
+    # in the output.) If you set :sudo_prompt to an empty string, Capissh
+    # will not send a preferred prompt.
+    def sudo_command(options={}, &block)
+      user = options[:as] && "-u #{options.delete(:as)}"
+
+      sudo_prompt_option = "-p '#{sudo_prompt}'" unless sudo_prompt.empty?
+      [configuration.fetch(:sudo, "sudo"), sudo_prompt_option, user].compact.join(" ")
+    end
+
+    # tests are too invasive right now
+    # protected
+
+    # Returns a Proc object that defines the behavior of the sudo
+    # callback. The returned Proc will defer to the +fallback+ argument
+    # (which should also be a Proc) for any output it does not
+    # explicitly handle.
+    def sudo_behavior_callback(fallback)
+      # in order to prevent _each host_ from prompting when the password
+      # was wrong, let's track which host prompted first and only allow
+      # subsequent prompts from that host.
+      prompt_host = nil
+
+      Proc.new do |ch, stream, out|
+        if out =~ /^Sorry, try again/
+          if prompt_host.nil? || prompt_host == ch[:server]
+            prompt_host = ch[:server]
+            logger.important out, "#{stream} :: #{ch[:server]}"
+            reset! :password
+          end
+        end
+
+        if out =~ /^#{Regexp.escape(sudo_prompt)}/
+          ch.send_data "#{configuration.fetch(:password,nil)}\n"
+        elsif fallback
+          fallback.call(ch, stream, out)
+        end
+      end
+    end
+
+    # Merges the various default command options into the options hash and
+    # returns the result. The default command options that are understand
+    # are:
+    #
+    # * :default_environment: If the :env key already exists, the :env
+    #   key is merged into default_environment and then added back into
+    #   options.
+    # * :default_shell: if the :shell key already exists, it will be used.
+    #   Otherwise, if the :default_shell key exists in the configuration,
+    #   it will be used. Otherwise, no :shell key is added.
+    def add_default_command_options(options)
+      defaults = configuration.fetch(:default_run_options, {})
+      options = defaults.merge(options)
+
+      env = configuration.fetch(:default_environment, {})
+      env = env.merge(options[:env]) if options[:env]
+      options[:env] = env unless env.empty?
+
+      shell = options[:shell] || configuration.fetch(:default_shell, nil)
+      options[:shell] = shell unless shell.nil?
+
+      options
+    end
+
+    # Returns the prompt text to use with sudo
+    def sudo_prompt
+      configuration.fetch(:sudo_prompt, "sudo password: ")
+    end
+
+    def continue_execution(tree)
+      if tree.branches.length == 1
+        continue_execution_for_branch(tree.branches.first)
+      else
+        tree.each { |branch| branch.skip! unless continue_execution_for_branch(branch) }
+        tree.any? { |branch| !branch.skip? }
+      end
+    end
+
+    def continue_execution_for_branch(branch)
+      case Capissh::CLI.debug_prompt(branch)
+      when "y"
+        true
+      when "n"
+        false
+      when "a"
+        exit(-1)
+      end
+    end
+
+  end
+end