#!/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')
%w{net-ssh net-scp sysinfo drydock}.each { |dir| $:.unshift File.join(File.dirname(__FILE__), '..', '..', dir, 'lib') }

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"
global :V, :version, "Display version number" do
  puts "Rye version: #{Rye::VERSION}"
  exit 0
end

before do |obj|
  # 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

end

option :u, :user, String, "User name to connect as"
option :i, :identity, String, "Private key"
argv :host
command :ssh do |obj|
  raise "No host supplied" if obj.argv.empty?
  opts = { :debug => true }
  opts[:user] = obj.option.user || Rye.sysinfo.user
  opts[:keys] = [obj.option.identity] || []
  obj.argv.each do |host|
    rbox = Rye::Box.new host, opts
    p rbox.connect
  end
end

about "Add your public keys to one or more remote machines"
usage "rye authorize [username@]host"
argv :hostname
command :authorize do |obj|
  raise "You must specify a host" unless obj.argv.hostname
  obj.argv.each do |hostname|
    if hostname.index('@')
      user, host = *hostname.split('@')
    else
      user, host = user = Rye.sysinfo.user, hostname
    end
    opts = { :debug => nil, :user => user, :auth_methods => %w(publickey hostbased) }
    puts "Authorizing #{user}@#{host}"
    rbox = Rye::Box.new(host, opts)
    rbox.authorize_keys_remote
    puts "Added public keys for: ", Rye.keys
    puts "Now try: " << "ssh #{user}@#{host}"
  end
  
end
command_alias :authorize, :authorise

about "Add your public keys to your current account on this machine"
command :authorize_local do |obj|
  user = Rye.sysinfo.user

  puts "Authorizing #{user}@localhost"
  rbox = Rye::Box.new('localhost', :debug => false)
  puts "Added public keys for: ", rbox.authorize_keys_local
  unless Rye.sysinfo.os == :windows
    puts "Now try: " << "ssh #{user}@localhost"
  end
end
command_alias :local_authorize, :local_authorise


about "Generate a public key from a private key"



about "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


about "Display your private keys"
command :keys do |obj|
  Rye.keys.each do |key|
    puts key
  end
end

about "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 :on

# We call 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 Drydock::ArgError, Drydock::OptError=> ex
  STDERR.puts ex.message
  STDERR.puts ex.usage
rescue => ex
  STDERR.puts "ERROR (#{ex.class.to_s}): #{ex.message}"
  STDERR.puts ex.backtrace if Drydock.debug?
rescue Interrupt
  puts "#{$/}Exiting... "
  exit 1
rescue SystemExit
  # Don't balk
end