ssh.rb

require 'delegate'
require 'net/ssh'
require 'net/sftp'

class SSH < Delegator
  class CommandFailed < RuntimeError
    attr_reader :command

    def initialize(command, message)
      @command = command
      super(message)
    end
  end

  attr_reader :connection

  def get_user(args)
    args.last.is_a?(Hash) && args.last.delete(:user) || ENV['USER'] || ENV['USERNAME'] || 'unknown'
  end

  def initialize(conn, *args)
    if conn.is_a?(Net::SSH::Connection::Session)
      @connection = conn
    else
      @connection = Net::SSH.start(conn, get_user(args), *args)
    end

    super(@connection)

    yield(self) if block_given?
  end

  def deinitialize
    
  end

  def __getobj__
    @connection
  end

  def __setobj__(obj)
    @connection = obj
  end

  def sftp
    @current_sftp ||= @connection.sftp.connect
  end

  def upload(source, target, options = {})
    sftp.upload!(source, target, options)
  end

  def download(source, target, options = {})
    sftp.download!(source, target, options)
  end

  def cat(source, options = {})
    sftp.download!(source, nil, options)
  end

  def write(target, data, options = {})
    unless data.is_a?(IO)
      data = StringIO.new(data)
    end

    sftp.upload!(data, target, options)
  end

  RUN_DEFAULTS = {
    stdin: false,
    stdout: true,
    stderr: true,
    raise: true,
  }

  def deferred_write(source, target)
    Thread.new do
      case source
      when IO
        while c = source.getc
          target.send_data c
        end
        target.eof!
      when Proc
        while c = source.call
          target.send_data c
        end
        target.eof!
      when nil, false
        target.eof!
      else
        target.send_data source.to_s
        target.eof!
      end
    end
  end

  EXEC_DEFAULTS = {
    stdout: $stdout,
    stderr: $stderr
  }

  def exec(cmd, options = {})
    run(cmd, EXEC_DEFAULTS.merge(options))
  end

  def run(cmd, options = {})
    options = RUN_DEFAULTS.merge(options)

    stdout = options[:stdout]
    stdout = StringIO.new if stdout == true

    stderr = options[:stderr]
    stderr = StringIO.new unless stderr.is_a?(IO)

    success = nil

    writer = nil

    channel = open_channel do |c, _|
      c.on_data {|_, d| stdout << d} if stdout
      c.on_extended_data {|_, t, d| stderr << d if t == 1}
      c.on_request('exit-status') {|_, d| success = d.read_long == 0}

      unless options[:stdin]
        c.eof!
      else
        writer = deferred_write(options[:stdin], c)
      end

      c.exec cmd
    end

    channel.wait

    ret = [success, stdout, stderr].each do |v|
      v.seek 0 if v.is_a?(StringIO)
    end

    if options[:raise] && !success
      message = stderr.is_a?(StringIO) ? stderr.read : 'see output'
      raise CommandFailed.new(cmd, message)
    end

    yield(stdout, stderr, success) if block_given?

    ret
  ensure
    if defined?(writer) && writer
      writer.kill if writer.alive?
      writer.join
    end
  end

  def subshell(script, method = :exec)
    __send__(method, "sh -c #{script.se}")
  end
end