#!/usr/bin/ruby

# Rye -- A CLI for some handy SSH tools
#
# If your reading this via the rdocs you won't be able to see the code
# See: http://github.com/delano/rye/blob/master/bin/rye
#
# Usage: rye
#

$:.unshift File.join(File.dirname(__FILE__), '..', 'lib')

require 'rubygems'
require 'stringio'
require 'yaml'
require 'drydock'
require 'rye'

include Drydock

global :p, :path, String, "A directory containing SSH private keys or the path to a private key"

before do |obj|
  STDERR.puts
  # Load private keys if specified
  if obj.global.path
    keys = Rye.find_private_keys(obj.global.path) 
    Rye.add_keys(keys)
  end
end

after do
  STDERR.puts
end

desc "Add your public keys to one or more remote machines"
option :u, :user, String, "Username"
argv :hostname
command :authorize do |obj|
  raise "You must specify a host" unless obj.argv.hostname
  
  highline = HighLine.new # Used for password prompt
  
  opts = { :debug => nil, :auth_methods => %w(publickey hostbased) }
  opts[:user] = obj.option.user if obj.option.user
  
  obj.argv.each do |hostname|
    
    
    retried = 0
    begin
      rbox = Rye::Box.new(hostname, opts).connect
      puts "Authorizing #{rbox.opts[:user]}@#{hostname}" if retried == 0
      
      # We know we're already authorized b/c we didn't have to give a password
      if retried == 0
        puts "%s is already authorized" % rbox.opts[:user]
        puts
        next
      end
      
    # An authentication failure means either the requested user doesn't
    # exist on the remote machine or we need to supply a password. 
    rescue Net::SSH::AuthenticationFailed => ex
      STDERR.puts
      retried += 1
      if retried <= 3
        #opts[:password] = highline.ask("Password: ") { |q| q.echo = '' }
        opts[:auth_methods] = %w(password keyboard-interactive)
        retry
      else
        STDERR.puts "Authentication failed."
        exit 1 
      end
    end
    
    puts "Added public keys for: ", rbox.authorize_keys
    
    puts
  end
  
end
command_alias :authorize, :authorise


desc "Generate a public key from a private key"



desc "Fetch the host keys for remote machines (suitable for your ~/.ssh/known_hosts file)"
usage "rye hostkey HOSTNAME [HOSTNAME2 HOSTNAME3 ...]"
argv :hostname
command :hostkeys do |obj|
  raise "You must specify a host" unless obj.argv.hostname
  ret = Rye.remote_host_keys(obj.argv.hostname)
  STDERR.puts $/, ret.stderr, $/
  puts ret.stdout
end
command_alias :hostkeys, :hostkey


desc "Display your private keys"
command :keys do |obj|
  Rye.keys.each do |key|
    puts key.join(' ')
  end
end

desc "Display your public SSH keys"
usage "rye pubkeys [--pem] [KEYPATH]"
option :p, :pem, "Output in PEM format"
argv :keypath
command :pubkeys do |obj|
  keys = obj.argv.empty? ? Rye.keys.collect { |k| k[2] } : obj.argv
  keys.each do |path|
    STDERR.puts "# Public SSH key for #{path}"
    k = Rye::Key.from_file(path)
    puts obj.option.pem ? k.public_key.to_pem : k.public_key.to_ssh2
  end
end
command_alias :pubkeys, :pubkey


default :keys
debug :off

# We can Drydock specifically otherwise it will run at_exit. Rye also
# uses at_exit for shutting down the ssh-agent. Ruby executes at_exit
# blocks in reverse order so if Drydock is required first, it's block
# will run after Rye shuts down the ssh-agent. 
begin
  Drydock.run!(ARGV, STDIN) if Drydock.run? && !Drydock.has_run?
rescue => ex
  STDERR.puts "ERROR (#{ex.class.to_s}): #{ex.message}"
  STDERR.puts ex.backtrace if Drydock.debug?
end