capissh 0.0.1

data/lib/capissh/logger.rb

@@ -0,0 +1,157 @@
+module Capissh
+  class Logger
+    attr_accessor :level, :device, :disable_formatters
+
+    IMPORTANT = 0
+    INFO      = 1
+    DEBUG     = 2
+    TRACE     = 3
+
+    MAX_LEVEL = 3
+
+    COLORS = {
+      :none     => "0",
+      :black    => "30",
+      :red      => "31",
+      :green    => "32",
+      :yellow   => "33",
+      :blue     => "34",
+      :magenta  => "35",
+      :cyan     => "36",
+      :white    => "37"
+    }
+
+    STYLES = {
+      :bright     => 1,
+      :dim        => 2,
+      :underscore => 4,
+      :blink      => 5,
+      :reverse    => 7,
+      :hidden     => 8
+    }
+
+    # Set up default formatters
+    @formatters = [
+      # TRACE
+      { :match => /command finished/,          :color => :white,   :style => :dim, :level => 3, :priority => -10 },
+      { :match => /executing locally/,         :color => :yellow,  :level => 3, :priority => -20 },
+
+      # DEBUG
+      { :match => /executing `.*/,             :color => :green,   :level => 2, :priority => -10, :timestamp => true },
+      { :match => /.*/,                        :color => :yellow,  :level => 2, :priority => -30 },
+
+      # INFO
+      { :match => /.*out\] (fatal:|ERROR:).*/, :color => :red,     :level => 1, :priority => -10 },
+      { :match => /Permission denied/,         :color => :red,     :level => 1, :priority => -20 },
+      { :match => /sh: .+: command not found/, :color => :magenta, :level => 1, :priority => -30 },
+
+      # IMPORTANT
+      { :match => /^err ::/,                   :color => :red,     :level => 0, :priority => -10 },
+      { :match => /.*/,                        :color => :blue,    :level => 0, :priority => -20 }
+    ]
+
+    class << self
+      def add_formatter(options)
+        @formatters.push(options)
+        @sorted_formatters = nil
+      end
+
+      def sorted_formatters
+        # Sort matchers in reverse order so we can break if we found a match.
+        @sorted_formatters ||= @formatters.sort_by { |i| -(i[:priority] || i[:prio] || 0) }
+      end
+
+      def default
+        @default ||= new(:level => 3)
+      end
+    end
+
+    def initialize(options={})
+      output = options[:output] || $stderr
+      if output.respond_to?(:puts)
+        @device = output
+      else
+        @device = File.open(output.to_str, "a")
+        @needs_close = true
+      end
+
+      @options = options
+      @level = options[:level] || 0
+      @disable_formatters = options[:disable_formatters]
+    end
+
+    def close
+      device.close if @needs_close
+    end
+
+    def log(level, message, line_prefix=nil)
+      if level <= self.level
+        # Only format output if device is a TTY and formatters are not disabled
+        if device.tty? && !@disable_formatters
+          color = :none
+          style = nil
+
+          Logger.sorted_formatters.each do |formatter|
+            if (formatter[:level] == level || formatter[:level].nil?)
+              if message =~ formatter[:match] || line_prefix =~ formatter[:match]
+                color = formatter[:color] if formatter[:color]
+                style = formatter[:style] || formatter[:attribute] # (support original cap colors)
+                message.gsub!(formatter[:match], formatter[:replace]) if formatter[:replace]
+                message = formatter[:prepend] + message unless formatter[:prepend].nil?
+                message = message + formatter[:append] unless formatter[:append].nil?
+                message = Time.now.strftime('%Y-%m-%d %T') + ' ' + message if formatter[:timestamp]
+                break unless formatter[:replace]
+              end
+            end
+          end
+
+          if color == :hide
+            # Don't do anything if color is set to :hide
+            return false
+          end
+
+          term_color = COLORS[color]
+          term_style = STYLES[style]
+
+          # Don't format message if no color or style
+          unless color == :none and style.nil?
+            unless line_prefix.nil?
+              line_prefix = format(line_prefix, term_color, term_style, nil)
+            end
+            message = format(message, term_color, term_style)
+          end
+        end
+
+        indent = "%*s" % [MAX_LEVEL, "*" * (MAX_LEVEL - level)]
+        (RUBY_VERSION >= "1.9" ? message.lines : message).each do |line|
+          if line_prefix
+            device.puts "#{indent} [#{line_prefix}] #{line.strip}\n"
+          else
+            device.puts "#{indent} #{line.strip}\n"
+          end
+        end
+      end
+    end
+
+    def important(message, line_prefix=nil)
+      log(IMPORTANT, message, line_prefix)
+    end
+
+    def info(message, line_prefix=nil)
+      log(INFO, message, line_prefix)
+    end
+
+    def debug(message, line_prefix=nil)
+      log(DEBUG, message, line_prefix)
+    end
+
+    def trace(message, line_prefix=nil)
+      log(TRACE, message, line_prefix)
+    end
+
+    def format(message, color, style, nl = "\n")
+      style = "#{style};" if style
+      "\e[#{style}#{color}m" + message.to_s.strip + "\e[0m#{nl}"
+    end
+  end
+end

data/lib/capissh/processable.rb

@@ -0,0 +1,54 @@
+module Capissh
+  module Processable
+    module SessionAssociation
+      def self.on(exception, session)
+        unless exception.respond_to?(:session)
+          exception.extend(self)
+          exception.session = session
+        end
+
+        return exception
+      end
+
+      attr_accessor :session
+    end
+
+    # perform ssh session looping with reading/writing until the block returns false
+    def process_iteration(wait=nil, &block)
+      ensure_each_session { |session| session.preprocess }
+
+      return false if block && !block.call(self)
+
+      readers = sessions.map { |session| session.listeners.keys }.flatten.reject { |io| io.closed? }
+      writers = readers.select { |io| io.respond_to?(:pending_write?) && io.pending_write? }
+
+      if readers.any? || writers.any?
+        readers, writers, = IO.select(readers, writers, nil, wait)
+      end
+
+      if readers
+        ensure_each_session do |session|
+          ios = session.listeners.keys
+          session.postprocess(ios & readers, ios & writers)
+        end
+      end
+
+      true
+    end
+
+    def ensure_each_session
+      errors = []
+
+      sessions.each do |session|
+        begin
+          yield session
+        rescue Exception => error
+          errors << SessionAssociation.on(error, session)
+        end
+      end
+
+      raise errors.first if errors.any?
+      sessions
+    end
+  end
+end

data/lib/capissh/server_definition.rb

@@ -0,0 +1,111 @@
+module Capissh
+  class ServerDefinition
+    include Comparable
+
+    # Wraps a string in a ServerDefinition, if it isn't already.
+    def self.wrap_server(item, options)
+      item.is_a?(ServerDefinition) ? item : ServerDefinition.new(item, options)
+    end
+
+    # Turns a list, or something resembling a list, into a properly-formatted
+    # ServerDefinition list. Keep an eye on this one -- it's entirely too
+    # magical for its own good. In particular, if ServerDefinition ever inherits
+    # from Array, this will break.
+    def self.wrap_list(*list)
+      options = list.last.is_a?(Hash) ? list.pop : {}
+      if list.length == 1
+        if list.first.nil?
+          return []
+        elsif list.first.is_a?(Array)
+          list = list.first
+        end
+      end
+      options.merge! list.pop if list.last.is_a?(Hash)
+      list.map do |item|
+        self.wrap_server item, options
+      end
+    end
+
+    # The default user name to use when a user name is not explicitly provided
+    def self.default_user
+      ENV['USER'] || ENV['USERNAME'] || "not-specified"
+    end
+
+    attr_reader :user, :host, :port, :options, :server
+
+    # Initialize a ServerDefinition with a string or object that describes the
+    # authority URI part, "user@host:port", for connecting with SSH.
+    #
+    # Any object that responds to the following methods, in order of priority,
+    # can be used in a ServerDefinition:
+    #
+    # 1. #host, #user (optional), and #port (optional, default: 22)
+    # 2. #authority - responding with something like "[user@]host[:port]"
+    # 3. #to_s      - responding with something like "[user@]host[:port]"
+    #
+    # If options are passed for the second argument, certain keys will be used:
+    #
+    # * :user - sets the user if one was not given in the authority
+    # * :port - sets the port if one was not given in the authority
+    # * :ssh_options - used for connecting with Net::SSH
+    def initialize(server, options={})
+      @server = server
+
+      if @server.respond_to?(:host)
+        @host = @server.host
+        @port = @server.port if @server.respond_to?(:port)
+        @user = @server.user if @server.respond_to?(:user)
+      else
+        if @server.respond_to?(:authority)
+          string = @server.authority
+        elsif @server.respond_to?(:to_s)
+          string = @server.to_s
+        else
+          raise ArgumentError, "Invalid server for ServerDefinition: #{@server.inspect}. Must respond to #host, #authority, or #to_s"
+        end
+        @user, @host, @port = string.match(/^(?:([^;,:=]+)@|)(.*?)(?::(\d+)|)$/)[1,3]
+      end
+
+      @options = options.dup
+      user_opt, port_opt = @options.delete(:user), @options.delete(:port)
+
+      @user ||= user_opt
+      @port ||= port_opt
+
+      @port = @port.to_i if @port
+    end
+
+    def <=>(server)
+      [host, port, user] <=> [server.host, server.port, server.user]
+    end
+
+    # Redefined, so that Array#uniq will work to remove duplicate server
+    # definitions, based solely on their host names.
+    def eql?(server)
+      host == server.host &&
+        user == server.user &&
+        port == server.port
+    end
+
+    alias :== :eql?
+
+    # Redefined, so that Array#uniq will work to remove duplicate server
+    # definitions, based on their connection information.
+    def hash
+      @hash ||= [host, user, port].hash
+    end
+
+    def connect_to_port
+      port || 22
+    end
+
+    def to_s
+      @to_s ||= begin
+        s = host
+        s = "#{user}@#{s}" if user
+        s = "#{s}:#{port}" if port && port != 22
+        s
+      end
+    end
+  end
+end

data/lib/capissh/ssh.rb

@@ -0,0 +1,110 @@
+require 'net/ssh'
+
+module Capissh
+  # A helper class for dealing with SSH connections.
+  class SSH
+    # Patch an accessor onto an SSH connection so that we can record the server
+    # definition object that defines the connection. This is useful because
+    # the gateway returns connections whose "host" is 127.0.0.1, instead of
+    # the host on the other side of the tunnel.
+    module Server
+      def self.apply_to(connection, server)
+        connection.extend(Server)
+        connection.xserver = server
+        connection
+      end
+
+      attr_accessor :xserver
+    end
+
+    # An abstraction to make it possible to connect to the server via public key
+    # without prompting for the password. If the public key authentication fails
+    # this will fall back to password authentication.
+    #
+    # +server+ must be an instance of ServerDefinition.
+    #
+    # If a block is given, the new session is yielded to it, otherwise the new
+    # session is returned.
+    #
+    # If an :ssh_options key exists in +options+, it is passed to the Net::SSH
+    # constructor. Values in +options+ are then merged into it, and any
+    # connection information in +server+ is added last, so that +server+ info
+    # takes precedence over +options+, which takes precendence over ssh_options.
+    def self.connect(server, options={})
+      connection_strategy(server, options) do |host, user, connection_options|
+        connection = Net::SSH.start(host, user, connection_options)
+        Server.apply_to(connection, server)
+      end
+    end
+
+    def self.gateway(gateway, options={})
+      connection_strategy(gateway, options) do |host, user, connection_options|
+        connection = Net::SSH::Gateway.new(host, user, connection_options)
+        Server.apply_to(connection, gateway)
+      end
+    end
+
+    # Abstracts the logic for establishing an SSH connection (which includes
+    # testing for connection failures and retrying with a password, and so forth,
+    # mostly made complicated because of the fact that some of these variables
+    # might be lazily evaluated and try to do something like prompt the user,
+    # which should only happen when absolutely necessary.
+    #
+    # This will yield the hostname, username, and a hash of connection options
+    # to the given block, which should return a new connection.
+    def self.connection_strategy(server, options={}, &block)
+      auth_methods = [ %w(publickey hostbased), %w(password keyboard-interactive) ]
+      password_value = nil
+
+      # construct the hash of ssh options that should be passed more-or-less
+      # directly to Net::SSH. This will be the general ssh options, merged with
+      # the server-specific ssh-options.
+      input_ssh_options = (options[:ssh_options] || {}).merge(server.options[:ssh_options] || {})
+
+      # load any SSH configuration files that were specified in the SSH options. This
+      # will load from ~/.ssh/config and /etc/ssh_config by default (see Net::SSH
+      # for details). Merge the explicitly given ssh_options over the top of the info
+      # from the config file.
+      ssh_options = Net::SSH.configuration_for(server.host, input_ssh_options.fetch(:config, true)).merge(input_ssh_options)
+
+      # Once we've loaded the config, we don't need Net::SSH to do it again.
+      ssh_options[:config] = false
+
+      ssh_options[:verbose] = :debug if options[:verbose] && options[:verbose] > 0
+
+      user =
+        server.user ||
+        options[:user] ||
+        ssh_options[:username] ||
+        ssh_options[:user] ||
+        ServerDefinition.default_user
+
+      port =
+        server.port ||
+        options[:port] ||
+        ssh_options[:port]
+
+      # the .ssh/config file might have changed the host-name on us
+      host = ssh_options.fetch(:host_name, server.host)
+
+      ssh_options[:port] = port if port
+
+      # delete these, since we've determined which username to use by this point
+      ssh_options.delete(:username)
+      ssh_options.delete(:user)
+
+      begin
+        connection_options = ssh_options.merge(
+          :password => password_value,
+          :auth_methods => ssh_options[:auth_methods] || auth_methods.shift
+        )
+
+        yield host, user, connection_options
+      rescue Net::SSH::AuthenticationFailed
+        raise if auth_methods.empty? || ssh_options[:auth_methods]
+        password_value = options[:password]
+        retry
+      end
+    end
+  end
+end

data/lib/capissh/transfer.rb

@@ -0,0 +1,218 @@
+require 'net/scp'
+require 'net/sftp'
+
+require 'capissh/processable'
+
+module Capissh
+  class Transfer
+    include Processable
+
+    def self.process(direction, from, to, sessions, options={}, &block)
+      new(direction, from, to, sessions, options, &block).process!
+    end
+
+    attr_reader :sessions
+    attr_reader :options
+    attr_reader :callback
+
+    attr_reader :transport
+    attr_reader :direction
+    attr_reader :from
+    attr_reader :to
+
+    attr_reader :logger
+    attr_reader :transfers
+
+    def initialize(direction, from, to, sessions, options={}, &block)
+      @direction = direction
+      @from      = from
+      @to        = to
+      @sessions  = sessions
+      @options   = options
+      @callback  = block
+
+      @transport = options.fetch(:via, :sftp)
+      @logger    = options.delete(:logger)
+
+      @session_map = {}
+
+      prepare_transfers
+    end
+
+    def process!
+      loop do
+        begin
+          break unless process_iteration { active? }
+        rescue Exception => error
+          if error.respond_to?(:session)
+            handle_error(error)
+          else
+            raise
+          end
+        end
+      end
+
+      failed = transfers.select { |txfr| txfr[:failed] }
+      if failed.any?
+        hosts = failed.map { |txfr| txfr[:server] }
+        errors = failed.map { |txfr| "#{txfr[:error]} (#{txfr[:error].message})" }.uniq.join(", ")
+        error = TransferError.new("#{operation} via #{transport} failed on #{hosts.join(',')}: #{errors}")
+        error.hosts = hosts
+
+        logger.important(error.message) if logger
+        raise error
+      end
+
+      logger.debug "#{transport} #{operation} complete" if logger
+      self
+    end
+
+    def active?
+      transfers.any? { |transfer| transfer.active? }
+    end
+
+    def operation
+      "#{direction}load"
+    end
+
+    def sanitized_from
+      if from.responds_to?(:read)
+        "#<#{from.class}>"
+      else
+        from
+      end
+    end
+
+    def sanitized_to
+      if to.responds_to?(:read)
+        "#<#{to.class}>"
+      else
+        to
+      end
+    end
+
+    private
+
+      def session_map
+        @session_map
+      end
+
+      def prepare_transfers
+        logger.info "#{transport} #{operation} #{from} -> #{to}" if logger
+
+        @transfers = sessions.map do |session|
+          session_from = normalize(from, session)
+          session_to   = normalize(to, session)
+
+          session_map[session] = case transport
+            when :sftp
+              prepare_sftp_transfer(session_from, session_to, session)
+            when :scp
+              prepare_scp_transfer(session_from, session_to, session)
+            else
+              raise ArgumentError, "unsupported transport type: #{transport.inspect}"
+            end
+        end
+      end
+
+      def prepare_scp_transfer(from, to, session)
+        real_callback = callback || Proc.new do |channel, name, sent, total|
+          logger.trace "[#{channel[:host]}] #{name}" if logger && sent == 0
+        end
+
+        channel = case direction
+          when :up
+            session.scp.upload(from, to, options, &real_callback)
+          when :down
+            session.scp.download(from, to, options, &real_callback)
+          else
+            raise ArgumentError, "unsupported transfer direction: #{direction.inspect}"
+          end
+
+        channel[:server] = session.xserver
+        channel[:host]   = session.xserver.host
+
+        return channel
+      end
+
+      class SFTPTransferWrapper
+        attr_reader :operation
+
+        def initialize(session, &callback)
+          session.sftp(false).connect do |sftp|
+            @operation = callback.call(sftp)
+          end
+        end
+
+        def active?
+          @operation.nil? || @operation.active?
+        end
+
+        def [](key)
+          @operation[key]
+        end
+
+        def []=(key, value)
+          @operation[key] = value
+        end
+
+        def abort!
+          @operation.abort!
+        end
+      end
+
+      def prepare_sftp_transfer(from, to, session)
+        SFTPTransferWrapper.new(session) do |sftp|
+          real_callback = Proc.new do |event, op, *args|
+            if callback
+              callback.call(event, op, *args)
+            elsif event == :open
+              logger.trace "[#{op[:host]}] #{args[0].remote}"
+            elsif event == :finish
+              logger.trace "[#{op[:host]}] done"
+            end
+          end
+
+          opts = options.dup
+          opts[:properties] = (opts[:properties] || {}).merge(
+            :server  => session.xserver,
+            :host    => session.xserver.host)
+
+          case direction
+          when :up
+            sftp.upload(from, to, opts, &real_callback)
+          when :down
+            sftp.download(from, to, opts, &real_callback)
+          else
+            raise ArgumentError, "unsupported transfer direction: #{direction.inspect}"
+          end
+        end
+      end
+
+      def normalize(argument, session)
+        if argument.is_a?(String)
+          argument.gsub(/\$CAPISSH:HOST\$/, session.xserver.host)
+        elsif argument.respond_to?(:read)
+          pos = argument.pos
+          clone = StringIO.new(argument.read)
+          clone.pos = argument.pos = pos
+          clone
+        else
+          argument
+        end
+      end
+
+      def handle_error(error)
+        raise error if error.message.include?('expected a file to upload')
+
+        transfer = session_map[error.session]
+        transfer[:error] = error
+        transfer[:failed] = true
+
+        case transport
+        when :sftp then transfer.abort!
+        when :scp  then transfer.close
+        end
+      end
+  end
+end