Launch AWS instances

launch_instances

#!/usr/bin/env ruby

# Require gems with versions
begin
  require 'logger'
  require 'rubygems'

  gem 'fog', '~> 0.6.0'
  gem 'trollop', '~> 1.16.2'
  require 'fog'; require 'trollop'
rescue LoadError
  $stderr.puts "Error: Requires rubygems ~> 1.3.7, fog ~> 0.6.0, trollop ~> 1.16.2"
  exit 1
end

# Setup Trollop options
opts = Trollop::options do
  version "launch_instances 2.0.3 (c) 2010 Josh Lindsey"
  banner <<-EOS
Launch one or several EC2 instances based on the set of [options].

Usage:
      launch_instances AMI [options]

where [options] are:
EOS

  opt :type, "instance type (size)", :default => 'm1.small'
  opt :num, "number of instances like this to launch", :default => 1
  opt :groups, "security groups to assign", :default => ['default']
  opt :key, "name of the keypair to use", :type => String
  opt :tags, "list of tag key/value pairs to apply", :type => :strings
  opt :file, "file to read in and assign to user_data", :type => :io
  opt :wait, "wait to return until instance is booted"
  opt :connect, "connect to instance after boot (implies -w)"
  opt :logfile, "file to log to", :type => :io, :default => $stdout
  opt :loglevel, "level to set the logger at (0-4; 0 == DEBUG)", :default => 1
end

ami = ARGV.pop

# Grab access keys from system vars
AccessKeyID = ENV['AWS_ACCESS_KEY_ID']
SecretAccessKey = ENV['AWS_SECRET_ACCESS_KEY']
if AccessKeyID.nil? or SecretAccessKey.nil?
  $stderr.puts "Error: must set the environment variables AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY"
  exit 1
end

# Input validation
AllowedInstanceTypes = %w(m1.small m1.large m1.xlarge t1.micro m2.xlarge m2.2xlarge
 x2.4xlarge c1.medium c1.xlarge cc1.4xlarge)
Trollop::die "AMI is required" if ami.nil?
Trollop::die :key, 'is required' if opts[:key].nil?
Trollop::die :type, 'must be a recognized AWS instance type 
  (See: http://aws.amazon.com/ec2/instance-types/)' unless AllowedInstanceTypes.include? opts[:type]
Trollop::die :tags, 'must be an even number (for key => val encoding into a hash)' if !opts[:tags].nil? and
  opts[:tags].length % 2 == 1
Trollop::die 'Cannot connect to more than one instance' if opts[:connect] and opts[:num] > 1
Trollop::die :loglevel, "must be between 0 and 4" unless (0..4).to_a.include? opts[:loglevel]

if opts[:connect] and !File.exists?(File.join(ENV['HOME'], '.ssh', "#{opts[:key]}.pem"))
  logger.warn "Key file #{opts[:key]}.pem not found in #{ENV['HOME']}/.ssh"
end

# Setup logger object
logger = Logger.new opts[:logfile]
logger.level = opts[:loglevel]

# Parse tags
tags = nil
unless opts[:tags].nil?
  tags = Hash[*opts[:tags]]
  logger.debug "Parsed tags: #{tags.inspect}"
end

# Init Fog instance
fog = Fog::Compute.new :provider => 'AWS',
                       :aws_access_key_id => AccessKeyID,
                       :aws_secret_access_key => SecretAccessKey

# Debug info
logger.debug "Options: #{opts.inspect}"
logger.debug "AMI: #{ami.inspect}"
logger.debug "Fog instance: #{fog.inspect}"

# Create the node
begin
  nodes = []
  logger.info "Launching #{opts[:num]} #{opts[:type]} node(s) from AMI '#{ami}'"
  
  opts[:num].times do
    node = fog.servers.create(:image_id => ami,
                              :flavor_id => opts[:type],
                              :groups => opts[:groups],
                              :key_name => opts[:key],
                              :user_data => (!opts[:file].nil? && opts[:file].read) || nil)

    unless tags.nil?
      logger.debug "Tagging node #{node.id}"
      tags.each_pair do |k, v|
        fog.tags.create :key => k, :value => v, :resource_id => node.id
      end
    end
    
    logger.debug "Node: #{node.reload.inspect}"
    nodes.push node
  end

rescue Fog::Service::NotFound => e
  if e.to_s =~ /AMI ID .* does not exist/
    logger.fatal "Provided AMI ID not found"
  else
    # Just re-emit the exception
    raise e
  end
  
  exit 1
rescue Fog::Service::Error => e
  if e.to_s =~ /AuthFailure/ 
    logger.fatal "Invalid AWS access credentials. Check your shell variables"
  elsif e.to_s =~ /requested instance type's architecture/
    logger.fatal "Provided instance size and AMI are incompatible"
  else
    # Just re-emit the exception
    raise e
  end

  exit 1
end

# Wait on the instance to boot
if opts[:wait] or opts[:connect]
  logger.info "Waiting for instance(s) to boot"
  
  loop do
    logger.debug "Sleeping"
    sleep 1
    break unless nodes.map { |node| node.reload.ready? }.include? false
  end
  
  # Connect to it if we're supposed to
  if opts[:connect]
    logger.info "Connecting to instance"
    exec %Q(ssh -o StrictHostKeyChecking=no -i #{File.join(ENV['HOME'], '.ssh', "#{opts[:key]}.pem")}\
      ubuntu@#{nodes.pop.reload.dns_name})
  end
  
  logger.info "DNS Names: "
  nodes.each { |node| logger.info node.dns_name }
else
  logger.info "Instance data: #{nodes.inspect}"
end

logger.info "Done"