#!/usr/bin/env ruby
#
# tn, copyright 2008 by Scott Severance <http://www.scottseverance.us/>
# Released under the terms of the GPL.
#
# This script makes it quick and easy to use SSH and SCP commands with multiple
# hosts, particularly if those hosts have long names or hostname and/or username
# details are difficult to remember or cumbersome to type.
# 
# Once configured, each hostname will have a user-chosen abbreviation, and one
# host can optionally be declared to be default, allowing it to be omitted
# altogether. Any additional options to be passed to ssh or scp can be specified
# and will be passed along. The configuration file allows for options to be
# passed by default to all hosts.
# 
# Examples:
# 
#   tn                   # Establishes an SSH connection to the default host
#   tn foo               # Establishes an SSH connection to the host with the
#                        # Abbreviation "foo".
#   tn scp foo __HOST__: # Uses SCP to transfer the file "foo" to the default
#                        # directory of the default host
# 
# For documentation, run tn --help. You'll first get instructions for
# setting up the config file. Once that's done, run tn --help again
# for usage instructions.
# 
# Depends on a system with OpenSSH installed. Might not work well outside of a
# UNIX-like environment. Tested in Ruby 1.8 and 1.9. Not tested in Ruby 2+.

$version = '0.2.1'

require 'yaml'

if ARGV[0] == '--debug'
  ARGV.shift
  $debug = true
  # Depending on the system, ruby-debug might be available directly, or it might
  # only be available after rubygems has been loaded. However, since rubygems is
  # not itself a dependency of tn, and is not necessarily installed on all
  # systems where ruby-debug is installed, it is only loaded if directly loading
  # ruby-debug fails.
  begin
    tries = -1
    begin
      tries += 1
      require 'ruby-debug'
    rescue LoadError
      raise if tries > 1
      require 'rubygems'
      retry
    end
    puts "Starting in debug mode.\n\nARGV = #{ARGV.inspect}\n\nType 'help' for a list of available commands. Type 'help <command>' for help on a particular command. To skip to the main program execution, type 'cont'."
    debugger
  rescue LoadError
    puts "ERROR: The debugger requires ruby-debug. To install the rubygems version of ruby-debug, type"
    puts "    sudo aptitude install rubygems && gem install ruby-debug"
    exit 3
  end
else $debug = false
end

class Main
  def initialize
    debugger if $debug
    if ARGV.include? '--config'
      file = nil
      found = false
      ARGV.each do |arg|
        if found
          file = arg
          break
        elsif arg == '--config'
          found = true
          next
        end
      end
      @config = ConfigFile.new(file).config
    else
      @config = ConfigFile.new.config
    end
    @@hosts = @config.fetch 'hosts'
    @@default_host_abbr = @config.fetch 'default'
    @@default_host = @@hosts[@@default_host_abbr]
    @args = []
    @host = nil
    @protocol = nil

    Main.usage if ARGV.include?('--help') || ARGV.include?('--version')
    alternate_config = false
    ARGV.each do |arg|
      if alternate_config
        alternate_config = false
        next
      elsif arg == '--config'
        alternate_config = true
        next
      elsif arg =~ /^(ssh|scp)$/
        raise ArgumentError, "Protocol #{arg} specified, but #{@protocol} was already given" if @protocol
        @protocol = arg
      elsif @@hosts.include? arg
        raise ArgumentError, "Host #{arg} specified, but #{@host} was already given" if @host
        @host = @@hosts[arg]
      else
        @args.push arg
      end
    end
    @protocol = 'ssh' unless @protocol
    @host = @@default_host unless @host
    @args.collect! do |arg|
      a = arg.gsub /__HOST__/, @host
      a = "'#{a}'" if a =~ /[^a-zA-Z0-9_\/.,@~:-]/
      a
    end
  end

  def process
    if @protocol == 'ssh'; process_ssh
    else process_scp
    end
  end

  def get_opts(which_type)
    begin
      opts = @config.fetch "#{which_type}_options"
      raise IndexError unless opts.respond_to?(:length) && opts.length >= 1
      opts.collect! { |o| "-o #{o}" }
      opts.push ''
    rescue IndexError
      opts = false
    end
    (opts) ? opts.join(' ') : ''
  end

  def process_ssh
    opts = get_opts 'ssh'
    str = "ssh #{opts}#{@host} #{@args.join ' '}"
    $stderr.puts "Executing: #{str}"
    exec str
  end

  def process_scp
    opts = get_opts 'scp'
    str = "scp #{opts}#{@args.join ' '}"
    $stderr.puts "Executing: #{str}"
    exec str
  end

  def self.usage(exit_status = 0)
    hosts = MyArray.new
    prog = File.split($0).last
    @@hosts.each_key { |k| hosts.push k }
    puts "#{prog} version #{$version}\n\n"
    puts "Usage: #{prog} [protocol] [host] [argument-1 [argument-2 [...]]]"
    puts "    protocol:        One of: ssh or scp."
    puts "                     Default: ssh\n\n"
    puts "    host:            One of: #{hosts.sort.join_sequence}"
    puts "                     Default: #{@@default_host_abbr}\n\n"
    puts "    other-arguments: Any additional arguments to pass (as understood by ssh or scp)\n\n"
    puts "Arguments may be passed in any order.  __HOST__ will substitute the proper"
    puts "host identifier. E.g., __HOST__:foo might result in user@somehost:foo.\n\n"
    puts "To configure the hosts #{prog} can connect to, edit #{ConfigFile.path}."
    puts "To specify an alternate config file, pass the option --config followed by the"
    puts "alternate filename.\n\n"
    puts "Debugging: #{prog} includes a built-in debugger. To run it, pass the argument"
    puts "--debug as the *first* argument."
    exit exit_status.to_i
  end
end

class MyArray < Array
  # Joins an array with a terminator. For example:
  #   array = MyArray[1, 2, 3]
  #   str = array.join_sequence
  #   puts str  # Result is: "1, 2, or 3"
  def join_sequence(terminator = 'or ', separator = ', ')
    push "%s%s" % [terminator, pop]
    join separator
  end
end

class ConfigFile
  attr_reader :path, :config

  def self.path
    @@path
  end

  def initialize(path = "#{ENV['HOME']}/.tn")
    @path = path
    @@path = @path
    begin
      config = File.open(@path) { |yaml| YAML::load yaml }
    rescue Errno::ENOENT
      File.open(@path, 'w') do |f|
        f.puts "#%YAML 1.1"
        f.puts "# This is the host configuration file for the tn program.\n\n"
        f.puts "hosts:"
        f.puts "  # Add hosts to this file following the example (MAKE SURE THEY'RE INDENTED)."
        f.puts "  # Field-specific notes:"
        f.puts "  # abbreviation: This is the short name you'll give on the command line."
        f.puts "  # username: This field is optional. If omitted, it will default to \"#{ENV['USER']}\".\n  #"
        f.puts "  # - abbreviation: some_abbreviation"
        f.puts "  #   username:     your_username"
        f.puts "  #   hostname:     your_full_hostname\n  #"
        f.puts "  # - abbreviation: another_abbreviation"
        f.puts "  #   username:     another_username"
        f.puts "  #   hostname:     another_hostname\n\n\n\n\n"
        f.puts "# Set the abbreviation of your default hostname here. (This should be identical"
        f.puts "# to one of your abbreviations listed above.)"
        f.puts "default: abbreviation_of_your_default_hostname\n\n"
        f.puts "# Specify any options to automatically pass to SSH using -o here. You can still"
        f.puts "# manually pass options when you invoke tn if you like."
        f.puts "ssh_options:"
        f.puts "  - ForwardX11=yes"
        f.puts "  - Compression=yes\n\n"
        f.puts "# As above, but for SCP."
        f.puts "scp_options:"
        f.puts "  - Compression=yes"
      end
      retry
    end

    if !config['hosts'].respond_to?(:length) || config['hosts'].length < 1
      raise NotConfiguredError
    end
    hosts = {}
    begin
      config['hosts'].each do |h|
        hst = h.fetch 'hostname'
        begin
          usr = h.fetch 'username'
          str = "%s@%s" % [usr, hst]
        rescue IndexError
          str = hst
        end
        hosts.store(h.fetch('abbreviation'), str)
      end
    rescue IndexError
      raise ConfigSyntaxError
    end
    @config = {'default' => config['default'],
               'hosts' => hosts,
               'ssh_options' => config['ssh_options'],
               'scp_options' => config['scp_options'],
              }
  end
end

class NotConfiguredError < StandardError
end

class ConfigSyntaxError < StandardError
end

if $0 == __FILE__
  begin
    Main.new.process
  rescue ArgumentError
    puts "ERROR: #{$!.exception}\n\n"
    Main.usage 1
  rescue NotConfiguredError
    print "#{File.split($0).last} hasn't been configured yet. You need to add one or more SSH/SCP"
    puts  " hosts to the file #{ConfigFile.path}. Once you do that, please try again."
    exit 2
  rescue ConfigSyntaxError
    print "ERROR: Your configuration file, #{ConfigFile.path}, contains a syntax error. Most likely this is the"
    print " result of a required field that is either missing or mis-typed. Please correct this"
    puts  " problem before continuing."
    exit 4
  end
end