Net::SSH::Session experiments

net_ssh_session.rb

require 'net/ssh'
require 'net/ssh/shell'
require 'open3'

module Net
  module SSH
    class SessionCommand
      attr_reader :command, :output, :exit_code

      # Initialize a new session command
      # @param command [String] original command
      # @param output [String] command execution output
      # @param exit_code [Integer] command execution exit code
      def initialize(command, output, exit_code)
        @command = command
        @output = output
        @exit_code = exit_code
      end

      # Check if exit code is successful
      # @return [Boolean]
      def success?
        exit_code == 0
      end

      # Check if exit code is not successful
      # @return [Boolean]
      def failure?
        exit_code != 0
      end

      # Get command output
      # @return [String]
      def to_s
        output
      end
    end
  end
end

module Net
  module SSH
    module SessionHelpers
      # Capture command output
      # @param command [String] command to execute
      # @return [String] command execution output
      def capture(command)
        run(command).output
      end

      # Check if remote file exists
      # @param path [String] remote file path
      # @return [Boolean]
      def file_exists?(path)
        result = capture("if [ -e #{path} ]; then echo 'true'; fi")
        result.strip == 'true'
      end

      # Check if remove process exists
      # @param pid [String] process id
      # @return [Boolean]
      def process_exists?(pid)
        result = capture("ps -p #{pid} ; true")
        result.strip.split("\n").size == 2
      end
    end
  end
end

module Net
  module SSH
    class Session
      include Net::SSH::SessionHelpers

      attr_reader :host, :user, :password
      attr_reader :connection, :shell
      attr_reader :options

      # Initialize a new ssh session
      # @param host [String] remote hostname or ip address
      # @param user [String] remote account username
      # @param password [String] remote account password
      def initialize(host, user, password='')
        @host = host
        @user = user
        @password = password
      end

      # Establish connection with remote server
      # @return [Boolean]
      def open
        @connection = Net::SSH.start(host, user, :password => password)
        @shell = @connection.shell
      end

      # Close connection with remote server
      # @return [Boolean]
      def close
        shell.close!
      end

      def on_output(&block)
        @on_output = block
      end

      # Execute command
      # @param command [String] command to execute
      # @param on_output [Block] output event block
      # @return [Integer] command execution exit code
      def exec(command, &on_output)
        status = nil
        shell.execute(command) do |process|
          process.on_output(&on_output)
          process.on_error_output(&on_output)
          process.on_finish { |p| status = p.exit_status }
        end
        shell.session.loop(1) { status.nil? }
        status
      end

      # Execute a single command
      # @param command [String] comand to execute
      # @return [SessionCommand]
      def run(command)
        output = ''

        exit_code = exec(command) do |process, data|
          output << data
        end

        SessionCommand.new(command, output, exit_code)
      end

      # Execute multiple commands
      # @param commands [Array] set of commands to execute
      # @param options [Hash] execution options
      # @return [Array] set of command execution results
      # 
      # Execution options are the following:
      # options[:break] - If set to `true`, execution chain will break on first failed command
      #
      def run_multiple(commands=[], options={})
        results = []

        [commands].flatten.compact.each do |cmd|
          results << run(cmd)
          break if results.last.failure? && options[:break] == true
        end

        results
      end
    end
  end
end