cloudflock 0.6.1 → 0.7.0

data/lib/cloudflock/remote/ssh.rb

@@ -1,311 +1,267 @@
 require 'cloudflock'
-require 'expectr'
+require 'timeout'
+require 'net/ssh'
 require 'socket'
+require 'thread'
 
-# Public: Wrap the tasks of logging into remote hosts via ssh and interacting
-# with them through Expectr.
-#
-# Examples
-#
-#   # Log in to remote host 'host.example.com'
-#   shell = SSH.new(host: 'host.example.com', pass: 'examplepass')
-class CloudFlock::Remote::SSH
-  # Public: Arguments to pass to ssh for when logging into remote hosts
-  SSH_ARGUMENTS = %w{-o UserKnownHostsFile=/dev/null -o
-  StrictHostKeyChecking=no -o NumberOfPasswordPrompts=1 -o
-  ConnectTimeout=15 -o ServerAliveInterval=30}.join(' ')
-
-  # Public: String used to standardize and identify the shell's prompt
-  PROMPT = '@@MIGRATE@@'
-
-  # Public: String representing the location of the history file
-  HISTFILE = '/root/.migration_history'
-
-  # Public: Initialize a new SSH object and automatically log in via SSH to
-  # the host with provided address/credentials.
+module CloudFlock; module Remote
+  # The SSH Class wraps the tasks of logging into a host and interacting with
+  # it via SSH.
   #
-  # args - A Hash used to specify optional arguments (default: {}):
-  #        :host         - String used to specify the address to which to
-  #                        connect.
-  #        :username     - String containing the remote user to use.
-  #                        (default: "root")
-  #        :password     - String containing the password to use for the user
-  #                        specified with :user. (default: "")
-  #        :port         - Fixnum specifying the port to which to connect.
-  #                        (default: 22)
-  #        :flush_buffer - Boolean specifying whether or not to flush output
-  #                        to STDOUT. (default: false)
-  #        :timeout      - Fixnum specifying the timeout for the Expectr object
-  #                        to use. (default: 30)
+  # Examples
   #
-  # Raises InvalidHostname if no host is specified.
-  # Raises InvalidHostname if looking up the host fails.
-  # Raises LoginFailed if logging in fails.
-  def initialize(args = {})
-    unless args[:host].kind_of?(String) && args[:host].length > 0
-      raise InvalidHostname, "No host specified"
+  #   # Log into root@host.example.com
+  #   shell = SSH.new(host: 'host.example.com', pass: 'examplepass')
+  #   shell.puts 'ls'
+  class SSH
+    # Public: String containing arguments to pass to ssh(1)
+    SSH_ARGUMENTS = '-o UserKnownHostsFile=/dev/null ' \
+                    '-o StrictHostKeyChecking=no '     \
+                    '-o NumberOfPasswordPrompts=1 '    \
+                    '-o ConnectTimeout=15 '            \
+                    '-o ServerAliveInterval=30'
+
+    # Public: Prompt to be set on a host upon successful login.
+    PROMPT = '@@CLOUDFLOCK@@'
+
+    # Public: Absolute path to the history file to use.
+    HISTFILE = '/root/.cloudflock_history'
+
+    # Internal: Default arguments to be used for SSH session initialization.
+    DEFAULT_ARGS = {username: '', password: '', ssh_key: '', port: 22}
+
+    # Internal: Default settings for calls to #batch.
+    DEFAULT_BATCH_ARGS = { timeout: 30, recoverable: true }
+
+    attr_reader :options
+
+    # Public: Create a new SSH object and log in to the specified host via ssh.
+    #
+    # args - Hash containing arguments relating to the SSH session. (default
+    #        defined in DEFAULT_ARGS):
+    #        :hostname       - String containing the address of the remote
+    #                          host.
+    #        :username       - String containing the remote user with which to
+    #                          log in.  (default: '')
+    #        :password       - String containing the password with which to log
+    #                          in.  (default: '')
+    #        :port           - Fixnum specifying the port to which to connect.
+    #                          (default: 22)
+    #        :ssh_key        - String containing the path to an ssh private
+    #                          key.  (default: '')
+    #        :key_passphrase - The passphrase for the ssh key if applicable.
+    #                          (default: '')
+    #
+    # Raises InvalidHostname if host lookup fails.
+    # Raises LoginFailed if logging into the host fails.
+    def initialize(args = {})
+      @options = sanitize_arguments(DEFAULT_ARGS.merge(args))
+      start_session
     end
 
-    begin
-      args[:host] = lookup_hostname(args[:host])
-    rescue SocketError
-      raise InvalidHostname, "Unable to look up host: #{args[:host]}"
+    # Public: Return the hostname of the host.
+    #
+    # Returns a String.
+    def hostname
+      options[:hostname]
     end
 
-    # Set up the rest of the arguments Hash
-    args[:username] ||= ''
-    args[:password] ||= ''
-    args[:flush_buffer] = false if args[:flush_buffer].nil?
-    args[:username] ||= 'root'
-    args[:timeout] ||= 30
-    args[:port] ||= 22
-    args[:public_key] ||= ''
-
-    # Sanitize the arguments Hash
-    args[:username].downcase!
-    args[:username].gsub!(/[^a-z0-9_-]/, '')
-    args[:port] = args[:port].to_i
-    args[:password].tr!("\u0000-\u001f\u007f\u2028-\u2029", '')
-
-    # Build the SSH command to send to the system
-    command = "ssh #{SSH_ARGUMENTS}"
-    if File.file?(File.expand_path(args[:public_key]))
-      command << " -i #{File.expand_path(args[:public_key])}"
-    end
-    if args[:username].length > 0
-      args[:username] << '@'
-    end
-    command << " #{args[:username]}#{args[:host]} -p #{args[:port]}"
-    @expect = Expectr.new(command, flush_buffer: args[:flush_buffer],
-                          timeout: args[:timeout])
+    # Public: Open a channel and execute an arbitrary command, returning any
+    # data returned over the channel.
+    #
+    # command     - Command to be executed.
+    # timeout     - Number of seconds to allow the command to run before
+    #               terminating the channel and returning any buffer returned
+    #               so far.  A value of 0 or nil will result in no timeout.
+    #               (default: 30)
+    # recoverable - Whether a Timeout should be considered acceptable.
+    #               (default: false)
+    # send_data   - Array containing data to be sent to across the channel
+    #               after the command has been run. (default: [])
+    #
+    # Returns a String.
+    # Raises Timeout::Error if timeout is reached.
+    def query(command, timeout = 30, recoverable = false, send_data = [])
+      buffer  = ''
+      running = 0
+
+      channel = @ssh.open_channel do |channel|
+        channel.request_pty
+
+        channel.exec(command) do |ch, success|
+          ch.on_data          { |_,    data| buffer << data }
+          ch.on_extended_data { |_, _, data| buffer << data }
+
+          ch.send_data(send_data.join + "\n")
+        end
+      end
 
-    raise LoginFailed unless login(args[:password])
-  rescue Timeout::Error, Expectr::ProcessError
-    raise LoginFailed
-  end
+      Timeout::timeout(timeout) { channel.wait }
 
-  # Public: Verify the host passed and return network address to which it's
-  # mapped.
-  #
-  # host - String containing the hostname or IP to verify
-  #
-  # Returns a String containing an IP
-  def lookup_hostname(host)
-    Socket.getaddrinfo(host, nil, nil, Socket::SOCK_STREAM)[0][3]
-  end
+      buffer.strip
+    rescue Timeout::Error
+      raise unless recoverable
+      channel.close
+      buffer.strip
+    rescue EOFError
+      start_session
+      retry
+    end
 
-  # Public: Wrap authentication and provide passwords if requested.  Upon
-  # detecting successful authentication, set the shell's PS1 to PROMPT.
-  #
-  # password - String containing the password to provide if requested.
-  #            (default: "")
-  #
-  # Returns true or false indicating login success.
-  def login(password = '')
-    @expect.expect(/password/i, true) do |match|
-      @expect.puts(password) if match.to_s =~ /password/i
+    # Public: Call query on a list of commands, allowing optional timeout and
+    # recoverable settings per command.
+    #
+    # commands - Array containing Hashes containing at minimum a command key.
+    #            Hash should follow the following specification:
+    #            command     - Command to execute.
+    #            timeout     - Timeout to specify for the call to #query.
+    #                          (default: 30)
+    #            recoverable - Whether the call should be considered
+    #                          recoverable if the timeout is reached.
+    #                          (default: true)
+    #
+    # Returns an Array containing results of each command.
+    def batch(commands)
+      commands.map! { |c| DEFAULT_BATCH_ARGS.merge(c) }
+      commands.map  { |c| query(c[:command], c[:timeout], c[:recoverable]) }
     end
 
-    # Wait to get a response (e.g. prompt), then set PS1
-    count = 0
-    continue = false
-    until continue
-      if count == 5
-        @expect.kill!
-        return false
+    # Public: Wrap query, guaranteeing that the user performing the given
+    # command is root.
+    #
+    # command     - Command to be executed.
+    # timeout     - Number of seconds to allow the command to run before
+    #               terminating the channel and returning any buffer returned
+    #               so far.  A value of 0 or nil will result in no timeout.
+    #               (default: 30)
+    # recoverable - Whether a Timeout should be considered acceptable.
+    #               (default: false)
+    #
+    # Returns a String.
+    # Raises Timeout::Error if timeout is reached.
+    def as_root(command, timeout = 30, recoverable = false)
+      return query(command, timeout, recoverable) if root?
+
+      priv    = 'su -'
+      priv    = 'sudo ' + priv if options[:sudo]
+      uid     = ['id;logout||exit']
+      command = ["#{command};logout||exit"]
+
+      passwordless = query(priv, timeout, true, uid)
+
+      unless /uid=0/.match(passwordless)
+        command.unshift(options[:root_password] + "\n")
       end
 
-      sleep 5
-      @expect.clear_buffer!
-      @expect.puts
-      continue = @expect.expect(/./, true).to_s =~ /./
-      count += 1
+      buffer = query(priv, timeout, recoverable, command)
+      cleanup = Regexp.new("^#{Regexp::escape(command.join)}\r\n")
+      buffer.gsub(cleanup, '')
     end
-    return false if @expect.pid == 0
-    @expect.puts("export PS1='#{PROMPT} '")
-    true
-  end
-
-  # Public: Check to see if the shell attached to the SSH object has superuser
-  # priveleges.
-  #
-  # Returns false if root privileges are detected, true otherwise.
-  def check_root
-    @expect.clear_buffer!
-    uid = query("UID_CHECK", command = "id")
-    root = /uid=0\(.*$/.match(uid)
-    root.nil?
-  end
 
-  # Public: Determine if currently logged in user is the superuser and, if
-  # not, attempt to gain superuser permissions via su/sudo.
-  #
-  # password - String containing password to use.
-  # use_sudo - Boolean value denoting whether to use sudo to obtain root
-  #            privileges. (default: false)
-  #
-  # Returns nothing.
-  # Raises StandardError if we are unable to obtain root.
-  def get_root(password, use_sudo = false)
-    match = nil
-    if check_root
-      @expect.send("sudo ") if use_sudo
-      @expect.puts("su -")
-      login(password)
+    # Public: Terminate the active ssh session.
+    #
+    # Returns nothing.
+    def logout!
+      @ssh.close
+      @ssh = nil
     end
 
-    raise RootFailed, "Unable to obtain root permissions" if check_root
+    private
 
-    @expect.puts("export PS1='#{PROMPT} '")
-    @expect.puts("export HISTFILE='#{HISTFILE}'")
-
-    set_timeout(5) do
-      while prompt(true)
-      end
+    # Private: Indicates whether the logged-in user is root.
+    #
+    # Returns true if the current user is UID 0, false otherwise.
+    def root?
+      @uid || @uid = fetch_uid
     end
-  end
 
-  # Public: Log out of any active shells.  Attempt a maximum of 10 logouts,
-  # then kill the ssh process if the Expectr object still reports a live pid.
-  #
-  # Returns nothing.
-  def logout!
-    count = 0
-    while @expect.pid > 0 && count < 10
-      count += 1
-      @expect.clear_buffer!
-      @expect.puts
-      prompt(true)
-      @expect.puts("exit")
-
-      sleep 1
+    # Private: Fetch the uid of the logged in user, return true if the uid is
+    # zero.
+    #
+    # Returns true if the current user is UID 0, false otherwise.
+    def fetch_uid
+      uid = query('id')
+      return false unless uid.gsub!(/.*uid=(\d+).*/, '\\1').to_i == 0
+      true
     end
 
-    @expect.kill!
-  rescue ArgumentError, Expectr::ProcessError
-    # Raised if puts or kill! fails.
-  end
+    # Internal: Sanitize arguments to be used
+    #
+    # args - Hash containing arguments relating to the SSH session:
+    #        :host           - String containing the address of the remote
+    #                          host.
+    #        :username       - String containing the remote user with which to
+    #                          log in.  (default: '')
+    #        :password       - String containing the password with which to log
+    #                          in.  (default: '')
+    #        :port           - Fixnum specifying the port to which to connect.
+    #                          (default: 22)
+    #        :ssh_key        - String containing the path to an ssh private
+    #                          key.  (default: '')
+    #        :key_passphrase - The passphrase for the ssh key if applicable.
+    #                          (default: '')
+    #
+    # Returns a Hash containing sanitized arguments suitable for passing to
+    # Net::SSH.  Not that #filter_ssh_arguments should be called to guarantee
+    # inappropriate options are not passed to Net::SSH::start.
+    def sanitize_arguments(args)
+      # Resolve the host to an IP
+      args[:host] = lookup_hostname(args[:hostname])
+      args[:port] = args[:port].to_i
+
+      # Username should be lowercase, alphanumeric only.
+      args[:username].downcase!
+      args[:username].gsub!(/[^a-z0-9_-]/, '')
+
+      # Remove control characters from password
+      args[:password].tr!("\u0000-\u001f\u007f\u2028-\u2029", '')
+
+      # Read in ssh key data if able.
+      if File.file?(args[:ssh_key])
+        key_path = File.expand_path(args[:ssh_key].to_s)
+        args[:key_data] = File.read(key_path)
+      end
 
-  # Public: Wait for a prompt from the SSH object.
-  #
-  # recoverable - Boolean specifying whether the prompt is recoverable if no
-  #               match is found. (default: false)
-  #
-  # Returns nothing.
-  def prompt(recoverable=false)
-    unless recoverable.kind_of?(TrueClass) || recoverable.kind_of?(FalseClass)
-      raise ArgumentError, "Should specify true or false"
+      args
+    rescue Errno::EACCES
+      args
     end
 
-    @expect.expect(PROMPT, recoverable)
-  end
-
-  # Public: Set the Expectr object's timeout.
-  #
-  # timeout - Fixnum containing the number of seconds which the Expectr object
-  #           should use as its timeout value, either temporarily or until set
-  #           explicitly again.
-  #
-  # Returns nothing.
-  # Yields nothing.
-  # Raises ArgumentError if timeout is not a Fixnum.
-  def set_timeout(timeout)
-    unless timeout.kind_of?(Fixnum) && timeout > 0
-      raise ArgumentError, "Expected an integer greater than 0"
+    # Internal: Filter all arguments not suitable to be passed to
+    # Net::SSH::start.
+    #
+    # args - Hash containing arguments to filter.
+    #
+    # Returns a Hash with only the keys :port, :password, :key_data and
+    # :key_passphrase defined.
+    def filter_ssh_options(args)
+      valid_arguments = [:port, :password, :key_data, :key_passphrase]
+      args.select { |opt| valid_arguments.include? opt }
     end
 
-    if block_given?
-      old_timeout = @expect.timeout
-      @expect.timeout = timeout
-      result = yield
-      @expect.timeout = old_timeout
-      return result
-    else
-      @expect.timeout = timeout
+    # Internal: Resolve the hostname provided and return the network address to
+    # which it maps.
+    #
+    # host - String containing the hostname or IP to verify.
+    #
+    # Returns a String containing an IP.
+    #
+    # Raises InvalidHostname if address information cannot be obtained for
+    # the specified host.
+    def lookup_hostname(host)
+      Socket.getaddrinfo(host, nil, nil, Socket::SOCK_STREAM)[0][3]
+    rescue SocketError
+      raise(InvalidHostname, Errstr::INVALID_HOST % host)
     end
-  end
-
-  # Public: Print a tag to a new line, followed by the output of an arbitrary
-  # command, then the tag again.  Return everything between the two tags.
-  #
-  # tag         - String containing the tag to use to isolate command output.
-  #               This string must not be empty after being constrained to
-  #               /[a-zA-Z0-9_-]/.  (default: "SSH_TAG")
-  # command     - Command to send to the active ssh session.
-  # recoverable - Whether a timeout should be considered recoverable or fatal.
-  #               (default: false)
-  #
-  # Returns a MatchData object from the Expectr object.
-  # Raises ArgumentError if tag or command aren't Strings.
-  # Raises ArgumentError if the tag String is empty.
-  def query(tag = "SSH_TAG", command = "", recoverable = false)
-    raise ArgumentError unless command.kind_of?(String) && tag.kind_of?(String)
-
-    tag.gsub!(/[^a-zA-Z0-9_-]/, '')
-    raise ArgumentError, "Alphanumeric tag required" if tag.empty?
-
-    @expect.send('printf "' + tag + '\n";')
-    @expect.send(command.gsub(/[\r\n]/, ' '))
-    @expect.puts(';printf "\n' + tag + '\n";')
-
-    match = @expect.expect(/^#{tag}.*^#{tag}/m, recoverable)
-    prompt(recoverable)
-
-    match.to_s.gsub(/^#{tag}(.*)^#{tag}$/m, '\1').strip
-  end
-
-  # Public: Set whether or not the Expectr object will flush the internal
-  # buffer to STDOUT.
-  #
-  # flush - Boolean denoting whether to flush the buffer.
-  #
-  # Returns nothing.
-  def flush_buffer(flush)
-    @expect.flush_buffer = flush
-  end
-
-  # Public: Provide access to the Expectr object's output buffer.
-  #
-  # Returns a String containing the buffer.
-  def buffer
-    @expect.buffer
-  end
-
-  # Public: Wrap Expectr#send.
-  #
-  # command - String containing the data to send to the Expectr object.
-  #
-  # Returns nothing.
-  def send(command)
-    @expect.send(command)
-  end
 
-  # Public: Wrap Expectr#puts.
-  #
-  # command - String containing the data to send to the Expectr object.
-  #           (default: "")
-  #
-  # Returns nothing.
-  def puts(command = "")
-    @expect.puts(command)
-  end
-
-  # Public: Wrap Expectr#expect.
-  #
-  # args - Variable length list of arguments to send to Expectr#expect.
-  #
-  # Returns a MatchData object once a match is found if no block is given.
-  # Yields the MatchData object representing the match.
-  # Raises TypeError if something other than a String or Regexp is passed.
-  # Raises Timeout::Error if a match isn't found, unless recoverable.
-  def expect(*args)
-    @expect.expect(*args)
-  end
-
-  # Public: Wrap Expectr#clear_buffer!
-  #
-  # Returns nothing.
-  def clear
-    @expect.clear_buffer!
+    # Internal: Start an SSH session.
+    #
+    # Sets @ssh.
+    #
+    # Returns nothing.
+    def start_session
+      ssh_opts = filter_ssh_options(options)
+      @ssh = Net::SSH.start(options[:hostname], options[:username], ssh_opts)
+    end
   end
-end
+end; end