ztk 0.2.4 → 0.2.5

data/lib/ztk/rescue_retry.rb

@@ -32,24 +32,22 @@ module ZTK
     class << self
 
       def try(options={}, &block)
-        !block_given? and raise RescueRetryError, "You must supply a block!"
+        options = Base.build_config({
+          :tries => 1,
+          :on => Exception
+        }.merge(options))
+        options.logger.debug { "options(#{options.inspect})" }
 
-        options = { :logger => $logger, :tries => 1, :on => Exception }.merge(options)
-
-        logger = options[:logger]
-        tries = options[:tries]
-        on = options[:on]
-
-        logger and logger.debug { options.inspect }
+        !block_given? and Base.log_and_raise(options.logger, RescueRetryError, "You must supply a block!")
 
         begin
           return block.call
-        rescue on => e
-          if ((tries -= 1) > 0)
-            logger and logger.warn { "Caught #{e.inspect}, we will give it #{tries} more tr#{tries > 1 ? 'ies' : 'y'}." }
+        rescue options.on => e
+          if ((options.tries -= 1) > 0)
+            options.logger.warn { "Caught #{e.inspect}, we will give it #{options.tries} more tr#{options.tries > 1 ? 'ies' : 'y'}." }
             retry
           else
-            logger and logger.fatal { "Caught #{e.inspect}, sorry, we have to give up now." }
+            options.logger.fatal { "Caught #{e.inspect}, sorry, we have to give up now." }
             raise e
           end
         end

data/lib/ztk/spinner.rb

@@ -36,15 +36,22 @@ module ZTK
 
     class << self
 
-      def spin(stdout=STDOUT)
+      def spin(options={}, &block)
+        options = Base.build_config({
+          :sleep => 0.1
+        }.merge(options))
+        options.logger.debug { "options(#{options.inspect})" }
+
+        !block_given? and Base.log_and_raise(options.logger, SpinnerError, "You must supply a block!")
+
         charset = %w( | / - \\ )
         count = 0
         spinner = Thread.new do
           while count do
-            stdout.print(charset[(count += 1) % charset.length])
-            stdout.print("\b")
-            stdout.respond_to?(:flush) and stdout.flush
-            sleep(0.1)
+            options.stdout.print(charset[(count += 1) % charset.length])
+            options.stdout.print("\b")
+            options.stdout.respond_to?(:flush) and options.stdout.flush
+            sleep(options.sleep)
           end
         end
         yield.tap do

data/lib/ztk/ssh.rb

@@ -101,37 +101,50 @@ module ZTK
     #   authentication.
     # @option config [String, Array<String>] :proxy_keys A single or series of
     #   identity files to use for authentication with the proxy.
-    def initialize(config={})
+    def initialize(configuration={})
       super({
         :forward_agent => true,
         :compression => false,
-        :user_known_hosts_file => '/dev/null'
-      }.merge(config))
+        :user_known_hosts_file => '/dev/null',
+        :timeout => 60
+      }.merge(configuration))
+      config.logger.debug { "config(#{config.inspect})" }
     end
 
     def inspect
-      user_host = "#{@config.user}@#{@config.host_name}"
-      port = (@config.port ? ":#{@config.port}" : nil)
-      [user_host, port].compact.join
+      tags = Array.new
+
+      if config.proxy_host_name
+        proxy_user_host = "#{config.proxy_user}@#{config.proxy_host_name}"
+        proxy_port = (config.proxy_port ? ":#{config.proxy_port}" : nil)
+        tags << [proxy_user_host, proxy_port].compact.join
+        tags << " >>> "
+      end
+
+      user_host = "#{config.user}@#{config.host_name}"
+      port = (config.port ? ":#{config.port}" : nil)
+      tags << [user_host, port].compact.join
+
+      tags.join.strip
     end
 
     # Starts an SSH session.  Can also be used to get the Net::SSH object.
     #
     # Primarily used internally.
     def ssh
-      @ssh ||= Net::SSH.start(@config.host_name, @config.user, ssh_options)
+      @ssh ||= Net::SSH.start(config.host_name, config.user, ssh_options)
     end
 
     # Starts an SFTP session.  Can also be used to get the Net::SSH object.
     #
     # Primarily used internally.
     def sftp
-      @sftp ||= Net::SFTP.start(@config.host_name, @config.user, ssh_options)
+      @sftp ||= Net::SFTP.start(config.host_name, config.user, ssh_options)
     end
 
     # Close our session gracefully.
     def close
-      log(:debug) { "close" }
+      config.logger.debug { "close" }
       ssh and !ssh.closed? and ssh.close
     end
 
@@ -147,8 +160,8 @@ module ZTK
     #   end
     #   ssh.console
     def console
-      log(:info) { "console(#{console_command.inspect})" }
-      log(:debug) { "config(#{@config.inspect})" }
+      config.logger.debug { "config(#{config.inspect})" }
+      config.logger.info { "console(#{console_command.inspect})" }
 
       Kernel.exec(console_command)
     end
@@ -183,63 +196,89 @@ module ZTK
         "#{header}\n"
       end
 
-      log(:debug) { "config(#{@config.inspect})" }
-      log(:info) { "exec(#{command.inspect}, #{options.inspect})" }
+      options = OpenStruct.new({ :exit_code => 0, :silence => false }.merge(options))
 
-      options = OpenStruct.new({ :silence => false }.merge(options))
-      log(:debug) { "options(#{options.inspect})" }
+      config.logger.debug { "config(#{config.inspect})" }
+      config.logger.debug { "options(#{options.inspect})" }
+      config.logger.info { "exec(#{command.inspect}, #{options.inspect})" }
 
       output = ""
+      exit_code = nil
+      exit_signal = nil
       stdout_header = false
       stderr_header = false
 
-      ZTK::RescueRetry.try(:tries => 3, :on => EOFError) do
-        @ssh = Net::SSH.start(@config.host_name, @config.user, ssh_options)
-
-        channel = ssh.open_channel do |chan|
-          log(:debug) { "Channel opened." }
-          direct_log(:debug) { log_header("OPENED") }
-
-          chan.exec(command) do |ch, success|
-            raise SSHError, "Could not execute '#{command}'." unless success
+      begin
+        Timeout.timeout(config.timeout) do
+          ZTK::RescueRetry.try(:tries => 3, :on => EOFError) do
+            @ssh = Net::SSH.start(config.host_name, config.user, ssh_options)
+
+            channel = ssh.open_channel do |chan|
+              config.logger.debug { "Channel opened." }
+
+              direct_log(:debug) { log_header("COMMAND") }
+              direct_log(:debug) { "#{command}\n" }
+              direct_log(:debug) { log_header("OPENED") }
+
+              chan.exec(command) do |ch, success|
+                success or log_and_raise(SSHError, "Could not execute '#{command}'.")
+
+                ch.on_data do |c, data|
+                  if !stdout_header
+                    direct_log(:debug) { log_header("STDOUT") }
+                    stdout_header = true
+                    stderr_header = false
+                  end
+                  direct_log(:debug) { data }
+
+                  config.stdout.print(data) unless options.silence
+                  output += data
+                end
+
+                ch.on_extended_data do |c, type, data|
+                  if !stderr_header
+                    direct_log(:warn) { log_header("STDERR") }
+                    stderr_header = true
+                    stdout_header = false
+                  end
+                  direct_log(:warn) { data }
+
+                  config.stderr.print(data) unless options.silence
+                  output += data
+                end
+
+                ch.on_request("exit-status") do |ch, data|
+                  exit_code = data.read_long
+                end
+
+                ch.on_request("exit-signal") do |ch, data|
+                  exit_signal = data.read_long
+                end
+
+                ch.on_open_failed do |c, code, desc|
+                  config.logger.fatal { "Open failed! (#{code.inspect} - #{desc.inspect})" }
+                end
 
-            ch.on_data do |c, data|
-              if !stdout_header
-                direct_log(:debug) { log_header("STDOUT") }
-                stdout_header = true
-                stderr_header = false
               end
-              direct_log(:debug) { data }
-
-              @config.stdout.print(data) unless options.silence
-              output += data
-            end
-
-            ch.on_extended_data do |c, type, data|
-              if !stderr_header
-                direct_log(:warn) { log_header("STDERR") }
-                stderr_header = true
-                stdout_header = false
-              end
-              direct_log(:warn) { data }
-
-              @config.stderr.print(data) unless options.silence
-              output += data
-            end
-
-            ch.on_open_failed do |c, code, desc|
-              log(:fatal) { "Open failed! (#{code.inspect} - #{desc.inspect})" }
             end
+            channel.wait
 
+            direct_log(:debug) { log_header("CLOSED") }
+            config.logger.debug { "Channel closed." }
           end
         end
-        channel.wait
 
-        direct_log(:debug) { log_header("CLOSED") }
-        log(:debug) { "Channel closed." }
+      rescue Timeout::Error => e
+        direct_log(:debug) { log_header("TIMEOUT") }
+        log_and_raise(SSHError, "Session timed out after #{config.timeout} seconds!")
       end
 
-      OpenStruct.new(:output => output, :exit => $?)
+      config.logger.debug { "exit_code(#{exit_code}), exit_signal(#{exit_signal})" }
+
+      if (exit_code != options.exit_code)
+        log_and_raise(SSHError, "exec(#{command.inspect}, #{options.inspect}) failed! [#{exit_code}, #{exit_signal}]")
+      end
+      OpenStruct.new(:output => output, :exit_code => exit_code, :exit_signal => exit_signal)
     end
 
     # Uploads a local file to a remote host.
@@ -258,23 +297,23 @@ module ZTK
     #   remote = File.expand_path(File.join("/tmp", "id_rsa.pub"))
     #   ssh.upload(local, remote)
     def upload(local, remote)
-      log(:debug) { "config(#{@config.inspect})" }
-      log(:info) { "upload(#{local.inspect}, #{remote.inspect})" }
+      config.logger.debug { "config(#{config.inspect})" }
+      config.logger.info { "upload(#{local.inspect}, #{remote.inspect})" }
 
       ZTK::RescueRetry.try(:tries => 3, :on => EOFError) do
-        @sftp = Net::SFTP.start(@config.host_name, @config.user, ssh_options)
+        @sftp = Net::SFTP.start(config.host_name, config.user, ssh_options)
         sftp.upload!(local.to_s, remote.to_s) do |event, uploader, *args|
           case event
           when :open
-            log(:debug) { "upload(#{args[0].local} -> #{args[0].remote})" }
+            config.logger.debug { "upload(#{args[0].local} -> #{args[0].remote})" }
           when :close
-            log(:debug) { "close(#{args[0].remote})" }
+            config.logger.debug { "close(#{args[0].remote})" }
           when :mkdir
-            log(:debug) { "mkdir(#{args[0]})" }
+            config.logger.debug { "mkdir(#{args[0]})" }
           when :put
-            log(:debug) { "put(#{args[0].remote}, size #{args[2].size} bytes, offset #{args[1]})" }
+            config.logger.debug { "put(#{args[0].remote}, size #{args[2].size} bytes, offset #{args[1]})" }
           when :finish
-            log(:debug) { "finish" }
+            config.logger.debug { "finish" }
           end
         end
       end
@@ -298,23 +337,23 @@ module ZTK
     #   remote = File.expand_path(File.join(ENV["HOME"], ".ssh", "id_rsa.pub"))
     #   ssh.download(remote, local)
     def download(remote, local)
-      log(:debug) { "config(#{@config.inspect})" }
-      log(:info) { "download(#{remote.inspect}, #{local.inspect})" }
+      config.logger.debug { "config(#{config.inspect})" }
+      config.logger.info { "download(#{remote.inspect}, #{local.inspect})" }
 
       ZTK::RescueRetry.try(:tries => 3, :on => EOFError) do
-        @sftp = Net::SFTP.start(@config.host_name, @config.user, ssh_options)
+        @sftp = Net::SFTP.start(config.host_name, config.user, ssh_options)
         sftp.download!(remote.to_s, local.to_s) do |event, downloader, *args|
           case event
           when :open
-            log(:debug) { "download(#{args[0].remote} -> #{args[0].local})" }
+            config.logger.debug { "download(#{args[0].remote} -> #{args[0].local})" }
           when :close
-            log(:debug) { "close(#{args[0].local})" }
+            config.logger.debug { "close(#{args[0].local})" }
           when :mkdir
-            log(:debug) { "mkdir(#{args[0]})" }
+            config.logger.debug { "mkdir(#{args[0]})" }
           when :get
-            log(:debug) { "get(#{args[0].remote}, size #{args[2].size} bytes, offset #{args[1]})" }
+            config.logger.debug { "get(#{args[0].remote}, size #{args[2].size} bytes, offset #{args[1]})" }
           when :finish
-            log(:debug) { "finish" }
+            config.logger.debug { "finish" }
           end
         end
       end
@@ -334,28 +373,19 @@ module ZTK
       command << [ "-o", "StrictHostKeyChecking=no" ]
       command << [ "-o", "KeepAlive=yes" ]
       command << [ "-o", "ServerAliveInterval=60" ]
-      command << [ "-i", @config.keys ] if @config.keys
-      command << [ "-p", @config.port ] if @config.port
-      command << [ "-o", "ProxyCommand=\"#{proxy_command}\"" ] if @config.proxy_host_name
-      command << "#{@config.user}@#{@config.host_name}"
+      command << [ "-i", config.keys ] if config.keys
+      command << [ "-p", config.port ] if config.port
+      command << [ "-o", "ProxyCommand=\"#{proxy_command}\"" ] if config.proxy_host_name
+      command << "#{config.user}@#{config.host_name}"
       command = command.flatten.compact.join(" ")
-      log(:debug) { "console_command(#{command.inspect})" }
+      config.logger.debug { "console_command(#{command.inspect})" }
       command
     end
 
     # Builds our SSH proxy command.
     def proxy_command
-      if !@config.proxy_user
-        message = "You must specify an proxy user in order to SSH proxy."
-        log(:fatal) { message }
-        raise SSHError, message
-      end
-
-      if !@config.proxy_host_name
-        message = "You must specify an proxy host_name in order to SSH proxy."
-        log(:fatal) { message }
-        raise SSHError, message
-      end
+      !config.proxy_user and log_and_raise(SSHError, "You must specify an proxy user in order to SSH proxy.")
+      !config.proxy_host_name and log_and_raise(SSHError, "You must specify an proxy host_name in order to SSH proxy.")
 
       command = ["ssh"]
       command << [ "-q" ]
@@ -364,12 +394,12 @@ module ZTK
       command << [ "-o", "StrictHostKeyChecking=no" ]
       command << [ "-o", "KeepAlive=yes" ]
       command << [ "-o", "ServerAliveInterval=60" ]
-      command << [ "-i", @config.proxy_keys ] if @config.proxy_keys
-      command << [ "-p", @config.proxy_port ] if @config.proxy_port
-      command << "#{@config.proxy_user}@#{@config.proxy_host_name}"
+      command << [ "-i", config.proxy_keys ] if config.proxy_keys
+      command << [ "-p", config.proxy_port ] if config.proxy_port
+      command << "#{config.proxy_user}@#{config.proxy_host_name}"
       command << "nc %h %p"
       command = command.flatten.compact.join(" ")
-      log(:debug) { "proxy_command(#{command.inspect})" }
+      config.logger.debug { "proxy_command(#{command.inspect})" }
       command
     end
 
@@ -378,29 +408,29 @@ module ZTK
       options = {}
 
       # These are plainly documented on the Net::SSH config class.
-      options.merge!(:encryption => @config.encryption) if @config.encryption
-      options.merge!(:compression => @config.compression) if @config.compression
-      options.merge!(:compression_level => @config.compression_level) if @config.compression_level
-      options.merge!(:timeout => @config.timeout) if @config.timeout
-      options.merge!(:forward_agent => @config.forward_agent) if @config.forward_agent
-      options.merge!(:global_known_hosts_file => @config.global_known_hosts_file) if @config.global_known_hosts_file
-      options.merge!(:auth_methods => @config.auth_methods) if @config.auth_methods
-      options.merge!(:host_key => @config.host_key) if @config.host_key
-      options.merge!(:host_key_alias => @config.host_key_alias) if @config.host_key_alias
-      options.merge!(:host_name => @config.host_name) if @config.host_name
-      options.merge!(:keys => @config.keys) if @config.keys
-      options.merge!(:keys_only => @config.keys_only) if @config.keys_only
-      options.merge!(:hmac => @config.hmac) if @config.hmac
-      options.merge!(:port => @config.port) if @config.port
-      options.merge!(:proxy => Net::SSH::Proxy::Command.new(proxy_command)) if @config.proxy_host_name
-      options.merge!(:rekey_limit => @config.rekey_limit) if @config.rekey_limit
-      options.merge!(:user => @config.user) if @config.user
-      options.merge!(:user_known_hosts_file => @config.user_known_hosts_file) if @config.user_known_hosts_file
+      options.merge!(:encryption => config.encryption) if config.encryption
+      options.merge!(:compression => config.compression) if config.compression
+      options.merge!(:compression_level => config.compression_level) if config.compression_level
+      options.merge!(:timeout => config.timeout) if config.timeout
+      options.merge!(:forward_agent => config.forward_agent) if config.forward_agent
+      options.merge!(:global_known_hosts_file => config.global_known_hosts_file) if config.global_known_hosts_file
+      options.merge!(:auth_methods => config.auth_methods) if config.auth_methods
+      options.merge!(:host_key => config.host_key) if config.host_key
+      options.merge!(:host_key_alias => config.host_key_alias) if config.host_key_alias
+      options.merge!(:host_name => config.host_name) if config.host_name
+      options.merge!(:keys => config.keys) if config.keys
+      options.merge!(:keys_only => config.keys_only) if config.keys_only
+      options.merge!(:hmac => config.hmac) if config.hmac
+      options.merge!(:port => config.port) if config.port
+      options.merge!(:proxy => Net::SSH::Proxy::Command.new(proxy_command)) if config.proxy_host_name
+      options.merge!(:rekey_limit => config.rekey_limit) if config.rekey_limit
+      options.merge!(:user => config.user) if config.user
+      options.merge!(:user_known_hosts_file => config.user_known_hosts_file) if config.user_known_hosts_file
 
       # This is not plainly documented on the Net::SSH config class.
-      options.merge!(:password => @config.password) if @config.password
+      options.merge!(:password => config.password) if config.password
 
-      log(:debug) { "ssh_options(#{options.inspect})" }
+      config.logger.debug { "ssh_options(#{options.inspect})" }
       options
     end
 

data/lib/ztk/tcp_socket_check.rb

@@ -90,11 +90,12 @@ module ZTK
     # @option config [Integer] :timeout (5) Set the IO select timeout.
     # @option config [Integer] :wait (60) Set the amount of time before the wait
     #   method call will timeout.
-    def initialize(config={})
+    def initialize(configuration={})
       super({
         :timeout => 5,
         :wait => 60
-      }.merge(config))
+      }.merge(configuration))
+      config.logger.debug { "config(#{config.inspect})" }
     end
 
     # Check to see if socket on the host and port specified is ready.  This
@@ -104,30 +105,21 @@ module ZTK
     # @return [Boolean] Returns true or false depending on weither the socket
     #   is ready or not.
     def ready?
-      if @config.host.nil?
-        message = "You must supply a host!"
-        log(:fatal) { message }
-        raise TCPSocketCheckError, message
-      end
-
-      if @config.port.nil?
-        message = "You must supply a port!"
-        log(:fatal) { message }
-        raise TCPSocketCheckError, message
-      end
+      config.host.nil? and log_and_raise(TCPSocketCheckError, "You must supply a host!")
+      config.port.nil? and log_and_raise(TCPSocketCheckError, "You must supply a port!")
 
-      socket = TCPSocket.new(@config.host, @config.port)
+      socket = TCPSocket.new(config.host, config.port)
 
-      if @config.data.nil?
-        log(:debug) { "read(#{@config.host}:#{@config.port})" }
-        ((IO.select([socket], nil, nil, @config.timeout) && socket.gets) ? true : false)
+      if config.data.nil?
+        config.logger.debug { "read(#{config.host}:#{config.port})" }
+        ((IO.select([socket], nil, nil, config.timeout) && socket.gets) ? true : false)
       else
-        log(:debug) { "write(#{@config.host}:#{@config.port}, '#{@config.data}')" }
-        ((IO.select(nil, [socket], nil, @config.timeout) && socket.write(@config.data)) ? true : false)
+        config.logger.debug { "write(#{config.host}:#{config.port}, #{config.data.size} bytes)" }
+        ((IO.select(nil, [socket], nil, config.timeout) && socket.write(config.data)) ? true : false)
       end
 
     rescue Errno::ETIMEDOUT, Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::EHOSTUNREACH => e
-      log(:debug) { "#{@config.host}:#{@config.port} - #{e.message}" }
+      config.logger.debug { "#{config.host}:#{config.port} - #{e.message}" }
       false
     ensure
       (socket && socket.close)
@@ -140,16 +132,16 @@ module ZTK
     # @return [Boolean] Returns true or false depending on weither the socket
     #   became ready or not.
     def wait
-      log(:debug) { "waiting for socket to become available; timeout after #{@config.wait} seconds" }
-      Timeout.timeout(@config.wait) do
+      config.logger.debug { "Waiting for socket to become available; timeout after #{config.wait} seconds." }
+      Timeout.timeout(config.wait) do
         until ready?
-          log(:debug) { "sleeping 1 second" }
+          config.logger.debug { "Sleeping 1 second." }
           sleep(1)
         end
       end
       true
     rescue Timeout::Error => e
-      log(:warn) { "socket(#{@config.host}:#{@config.port}) timeout!" }
+      config.logger.warn { "socket(#{config.host}:#{config.port}) timeout!" }
       false
     end