require 'net/ssh'

module SSHCommands
  def execute(host, username, key_path, cmd)
    Net::SSH.start(host, username, :keys => [ key_path ]) do |ssh|
      command_output = ssh_exec!(ssh, cmd)
      stdout = command_output[0]
      stderr = command_output[1]
      exit_code = command_output[2]

      return {
        "stdout" => stdout,
        "stderr" => stderr,
        "exit_code" => exit_code
      }
    end
  end

  def ssh_exec!(ssh, command)
    # Originally submitted by 'flitzwald' over here: http://stackoverflow.com/a/3386375
    stdout_data = ""
    stderr_data = ""
    exit_code = nil

    ssh.open_channel do |channel|
      channel.exec(command) do |ch, success|
        unless success
          abort "FAILED: couldn't execute command (ssh.channel.exec)"
        end
        channel.on_data do |ch_, data|
          stdout_data += data
        end

        channel.on_extended_data do |ch_, type, data|
          stderr_data += data
        end

        channel.on_request("exit-status") do |ch_, data|
          exit_code = data.read_long
        end
      end
    end
    ssh.loop
    return [stdout_data, stderr_data, exit_code]
  end
end
