#!/usr/bin/ruby

require 'optparse'
require 'rubygems'
require 'aws-sdk'

options = {
  'use_public_ip'  => false,
  'region'         => 'us-west-2'
}

OptionParser.new do |opts|
  opts.banner = "Usage: ec2-i2ssh [-t TAGS] [-l USER] [-k KEY -s SECRET] [-r region]"

  opts.on("-l", "--login [USER]", "Log in with this user") do |opt|
    options['user'] = opt
  end

  opts.on("-t", "--tags [TAGS]", Array, "a 'comma' sparated key value pair of tags and values (i.e. role=web|database,environment=dev)") do |opt|
    options['tags'] = opt
  end

  opts.on("-r", "--region [REGION]", "AWS region") do |opt|
    options['region'] = opt
  end

  opts.on("-s", "--screen [SCREEN]", "What screen to use for clustering windows (form multiple displays)") do |opt|
    options['screen'] = opt
  end

  opts.on("-p", "--use-public-ip", "Use public IP (default #{options['use_public_ip']})") do
    options['use_public_ip'] = true
  end

end.parse!

Aws.config.update({:region => options['region']})
$ec2 = Aws::EC2::Client.new

def get_ip(instance, options)
  if options['use_public_ip']
    return instance.public_ip_address
  else
    return instance.private_ip_address
  end
end

def get_instances_by_tag (tag_key,tag_values,options)
  instances = []
  response = $ec2.describe_instances(filters: [{:name => "tag:#{tag_key}", :values => tag_values},{:name => "instance-state-name", :values => ['running']}])
  if response.reservations[0].nil?
    puts "No running instances in #{options['region']} had a tag named '#{tag_key}' with values '#{tag_values}'"
  else
    response.reservations.each do |reservation|
      reservation.instances.each do |instance|
        instances << get_ip(instance, options)
      end
    end
  end
  return instances
end

matched_instances = []
# filter instances based on tags
options['tags'].each do |tag|
  tag = tag.split('=')
  key = tag[0]
  values = tag[1].split('|')
  if matched_instances.empty?
    matched_instances = get_instances_by_tag(key,values,options)
  end
  matched_instances = matched_instances & get_instances_by_tag(key,values,options)
end

if matched_instances.length == 0
  puts "No instnaces matched filter: #{options['tags']}"
  exit 1
end

cssh = (/darwin/ =~ RUBY_PLATFORM) ? 'i2cssh' : 'cssh'

cmd = "#{cssh}"
cmd = cmd + " -l #{options['user']}" unless options['user'].nil?
cmd = cmd + " -screen #{options['screen']}" unless options['screen'].nil?
cmd = cmd + " #{matched_instances.join ' '}"

puts "Connecting to #{matched_instances.length} instances"
puts cmd
exec cmd
