SSH wrapper for working with screen & clusters

nssh.rb

#!/usr/bin/env ruby

# This script does minimal arg parsing to filter out SSH arguments to
# be mostly compatible with plain SSH command line syntax. It's not
# 100% but good enough.
#
# usage: nssh [-1246AaCfgKkMNnqsTtVvXxY] [-b bind_address] [-c cipher_spec]
#           [-D [bind_address:]port] [-e escape_char] [-F configfile]
#           [-i identity_file] [-L [bind_address:]port:host:hostport]
#           [-l login_name] [-m mac_spec] [-O ctl_cmd] [-o option] [-p port]
#           [-R [bind_address:]port:host:hostport] [-S ctl_path]
#           [-w local_tun[:remote_tun]] [user@]hostname
#
# nssh-specific arguments:
#   --list <machine_list>   A machine list is a dsh-style ~/.dsh/machines.$listname.

require 'resolv'

class NamedSSH
  attr_reader :dsh_config_dir
  attr_reader :dsh_config_file
  attr_reader :dsh_list
  attr_reader :nssh_last_file
  attr_reader :ssh_args
  attr_accessor :hostname

  def initialize(options = {})
    @ssh_args = options[:ssh_args] || Array.new
    @dsh_list = options[:dsh_list] || "machines.list"

    @dsh_config_dir  = options[:dsh_config_dir]
    @dsh_config_file = File.join(@dsh_config_dir, @dsh_list)
    @nssh_last_file  = options[:nssh_last_file]
    @hostname        = options[:hostname]

    unless @hostname != nil and @hostname.length > 3
      raise "Invalid hostname '#@hostname'."
    end
  end

  # most of the arg parsers looked painful to do what this does;
  # it needs to stash & ignore SSH options, while parsing out --list
  # and grab the hostname
  def self.parse_options
    ssh_args = Array.new
    dsh_list = nil
    hostname = nil

    # manual argument parsing - be intelligent about perserving ssh
    # options while adding custom options for nssh
    idx=0
    loop do
      break if idx == ARGV.size

      #puts "Arg #{idx}: #{ARGV[idx]}"

      # ssh switches
      if ARGV[idx].match(/^-[1246AaCfgKkMNnqsTtVvXxY]$/)
        ssh_args << ARGV[idx]
        #puts "ssh switch #{ssh_args[-1]}"

      # ssh options that take a value
      elsif ARGV[idx].match(/^-[bcDeFiLlmOopRSw]$/)
        ssh_args << ARGV[idx]
        idx+=1
        # force quoting - it should never hurt and makes stuff like -o options work correctly
        ssh_args << ARGV[idx]
        #puts "ssh args #{ssh_args[-2]} #{ssh_args[-1]}"

      # allow specification of a .dsh list in my style where --list foobar resolves to
      # ~/.dsh/machines.foobar to use with "nssh --list foobar next"
      elsif ARGV[idx].match(/^--list/)
        idx+=1
        dsh_list = "machines." << ARGV[idx]
        #puts "dsh list: #{dsh_list}"

      # user@hostname is a definite match
      # split it and use -u instead because hostname needs to be standalone
      elsif ARGV[idx].match(/^\w+@[-\.\w]+$/)
        user, hostname = ARGV[idx].split '@'
        ssh_args << '-o' << "User #{user}"
        #puts "user@hostname: user: #{user}, hostname: #{hostname}"

      # a bare, uncaptured argument is likely the hostname
      else
        hostname = ARGV[idx]
        #puts "hostname: #{hostname}"
      end

      idx+=1
    end

    return {
      :ssh_args => ssh_args,
      :dsh_list => dsh_list,
      :hostname => hostname
    }
  end

  # read the last host from "nssh next" iteration from a file
  def read_last()
    host = nil

    if File.exists?(@nssh_last_file)
      File.open(@nssh_last_file, 'r') do |f|
        host = f.gets
        host.chomp! if host != nil
      end
    end

    return host
  end

  # save the last host for "nssh next" iteration
  def save_last
    File.open(@nssh_last_file, 'w') do |f|
      f.puts @hostname
    end
  end
  # read the dsh machines file and return the next host in the list
  # after whatever was in the @nssh_last_file
  def read_next
    last = read_last()
    next_host = nil

    unless File.exists?(@dsh_config_file)
      raise "#@dsh_config_file does not exist on the filesystem. --list #@dsh_list is not valid."
    end

    File.open(@dsh_config_file, 'r') do |f|
      # last host is not defined, return first in file
      if last == nil
        next_host = f.gets.chomp
      end

      while (candidate = f.gets)
        candidate.chomp!
        if candidate == last
          if f.eof?
            raise "Reached end of #@dsh_config_file. There is no next host!"
          else
            next_host = f.gets.chomp
          end
        end
      end
    end

    return next_host
  end
end

# use class method to parse ARGV
options = NamedSSH.parse_options

nssh = NamedSSH.new(
  :dsh_config_dir => File.join(ENV['HOME'], ".dsh"),
  :nssh_last_file => File.join(ENV['HOME'], '.nssh-last'),
  :hostname       => options[:hostname],
  :ssh_args       => options[:ssh_args],
  :dsh_list       => options[:dsh_list]
)

# reset the position in the machine list to the top
if nssh.hostname == "reset"
  File.unlink(nssh.nssh_last_file)
  exit
end

# choose the next host in the machine list, great for firing  up
# a ton of screen windows in a row in an already-running screen
# If I'm logging into a whole cluster in an existing screen session, I'll load
# "nssh next --list $cluster" into my clipboard then ...
# ctrl-a n, <paste>, <enter>, ctrl-a n <paste> <enter>, ...
# (my screenrc spawns with 256 open & ready shells)
if nssh.hostname == "next"
  nssh.hostname = nssh.read_next
  nssh.save_last
end

# set the terminal title in GNU Screen
puts "\033k#{nssh.hostname}\033\\"

# set an environment variable to the selected hostname to
# pass through SSH. LC_* is accepted by default in most ssh servers
# this is useful for setting a PS1 with a meaningful CNAME on the remote
# host via .profile (esp handy for EC2 boxen)
# e.g.
# ps1host=$(hostname)
# [ -n "$LC_UI_HOSTNAME" ] && ps1host=$LC_UI_HOSTNAME
# if [[ ${EUID} == 0 ]] ; then
#   PS1="\\[\\033[01;31m\\]$ps1host\\[\\033[01;34m\\] \\W \\$\\[\\033[00m\\] "
# else   
#   PS1="\\[\\033[01;32m\\]\\u@$ps1host\\[\\033[01;34m\\] \\w \\$\\[\\033[00m\\] "
# fi
ENV['LC_UI_HOSTNAME'] = nssh.hostname

# resolve hostnames to IP then back to the IP's name because
# I don't see where resolv lets you just get the CNAME
# this helps for stuff where we use CNAME's to point at e.g. EC2 hosts
# to switch to the EC2 name before ssh'ing but still set all the display
# stuff to the human hostname
#
# I use this to make ssh match config entries in ~/.ssh/config which is generated
# based on the EC2 instance name. This saves me having to do backflips to provide
# extra aliases in ssh_config while still letting me nssh to the CNAME's.
resolver = Resolv.new
addr = resolver.getaddress nssh.hostname
realname = resolver.getname addr

# run SSH
Kernel.exec("ssh", *nssh.ssh_args, realname)