ssh-with-lookup

#!/usr/bin/env ruby
require 'yaml'

SERVERS = YAML::load_file(File.join(ENV['HOME'], "/etc/ssh-with-lookup.yml"), aliases: true)

target = ARGV.shift
raise ArgumentError, "usage: ssh-with-lookup.rb <host>.<key> [ssh params ... ]" if target.nil?

*host, key = target.split('.')
rest       = ARGV.join(' ')
host       = host.join('.')
raise ArgumentError, "usage: ssh-with-lookup.rb <host>.<key> [ssh params ... ]" if host.nil? || key.nil?
raise ArgumentError, "unknown key" unless SERVERS.has_key?(key)

# grab the default config for this key
conf = SERVERS[key]['default']

# if there's a specific host entry, override from that
if SERVERS[key].has_key?(host)
  conf.merge!(SERVERS[key][host])
elsif conf.has_key?('host-suffix') # for defaulting hosts, append the key-wide host suffix if supplied
  conf['host'] = host + ".#{conf['host-suffix']}"
end
raise ArgumentError, "unknown host" unless conf.has_key?('host')
conf['options'] ||= ''

# if `-l <user>` was on the commandline, move it to conf['options']
if user_opt = rest[/-l\s*\w+/]
  rest.gsub!(user_opt, '')
  conf['options'].gsub!(/-l\s*\w+/, '')
  conf['options'] += " #{user_opt}"
end

# discover the file containing the ssh-agent details we want
agentenv = ENV['HOME'] + "/.ssh/.ssh-agent-#{conf['key']}"
raise RuntimeError, "no agent config found at #{agentenv}" unless File.exist?(agentenv)

# set this process's environment from the file's contents
File.open(agentenv, 'r') do |f|
  while line = f.gets do
    line.gsub!(/;.*/, '')
    next unless line.match(/^SSH/)
    key, value = line.chomp.split('=')
    ENV[key] = value
  end
end

# do it
case
when $0.match(/ssh/)
  cmd = "ssh #{conf['options']} #{conf['host']} #{rest}"
  system("tmux rename-window -t #{ENV['TMUX_PANE']} #{host}") if ENV.has_key?('TMUX')
when $0.match(/scp/)
  # sigh, why would they use the same flags
  conf['options'].gsub!(/-l\s*/, '-oUser=')
  conf['options'].gsub!(/-A/, '')
  conf['options'].gsub!(/-p/, '-P')

  rest.gsub!(/:/, "#{conf['host']}:")
  cmd = "scp #{conf['options']} #{rest}"
end

ENV['SSH_ASKPASS'] = ENV['HOME'] + '/bin/ssh-askpass'
STDERR.puts cmd
exec cmd