Amazon VPC /etc/hosts hack

vpc-hosts-hack.rb

#!/usr/bin/env ruby

require 'pp'
require 'json'
require 'fileutils'
require 'set'


# Filter for specific Route 53 hosted zones
#
Relevant_Zones = ['mbst.tv.','metabroadcast.com.']


# First, gather all CNAME records
#   { target name => [referencing names] }
#
$hostnames = Hash.new {|h,k| h[k] = []}

zones = JSON.parse(`aws route53 list-hosted-zones`)['HostedZones']
zones.each do |z|
  unless Relevant_Zones.include? z['Name']
    puts "Skipping #{z['Name']}"
    next
  end
  zoneId = z['Id']
  rsets = JSON.parse(`aws route53 list-resource-record-sets --hosted-zone-id #{zoneId}`)['ResourceRecordSets']
  (rsets || []).each do |rset|
    target = rset['ResourceRecords'][0]['Value'] rescue next
    $hostnames[target] << rset['Name'][0..-2] rescue next  # slice is to remove trailing dots
  end
end


# Function that will get the complete set of names pointing at a target
#   from the graph represented by $hostnames
#
def allHostnameAliases(hostname, _addToSet = Set.new)
  _addToSet << hostname
  $hostnames[hostname].each do |h|
    allHostnameAliases(h, _addToSet)
  end
  _addToSet
end


# Next, fetch a list of running instances
#   { instance-id => { details… } }
#
instances = Hash.new {|h,k| h[k] = {id:k}}

JSON.parse(`aws ec2 describe-instances`)['Reservations'].each do |r|
  r['Instances'].each do |i|
    id = i['InstanceId']
    if i['PrivateIpAddress'] && i['PublicDnsName']
      instances[id][:extDns] = i['PublicDnsName']
      instances[id][:hostnames] = $hostnames["#{i['PublicDnsName']}."]
      instances[id][:intIp] = i['PrivateIpAddress']
    end
  end
end


# Now map every aliasing hostname for an instance to its internal IP
#
mappings = {}

vpcInstances = instances.select {|k,v| v[:intIp].start_with?('172.')}  # IP range of instances in VPC
vpcInstances.values.each do |i|
  if i[:hostnames].empty?
    puts "Hostname not known for IP #{i[:intIp]}"
    next
  end
  i[:hostnames].each do |h|
    allHostnameAliases(h).each do |hh|
      mappings[hh] = i[:intIp]
    end
  end
end


# Update the hosts file with 'mappings'
#
HOSTSFILE = '/etc/hosts'
HOSTSFILE_TMP = '/tmp/hosts.vpc-hosts-hack'  # Write here, to not immediately need sudo

FLINE = '# BEGIN mbst vpc fixes'  # Marker lines bounding the section of /etc/hosts that
LLINE = '# END mbst vpc fixes'    #   this script will replace or add

if ARGV[0] == 'dump'
  lines = []  # Current hosts file
  File.open(HOSTSFILE, 'r') do |f|
    lines = f.each_line.entries.map {|l| l.strip}
  end

  newLines = [FLINE, LLINE]
  newLines[1,0] = mappings.map {|k,v| "#{v}\t#{k}" }

  # Does our section already exist in the current file?
  #
  first = lines.find_index FLINE
  last = lines.find_index LLINE

  if (first.nil? and last.nil?)
    lines.concat newLines             # Create our section
  elsif not (first.nil? or last.nil?)
    lines[first..last] = newLines     # Replace existing section
  else
    raise 'WTF, the hosts file is mangled'
  end

  File.open("#{HOSTSFILE_TMP}.new", 'w') do |f|
    f << lines.join("\n")
  end
  FileUtils.cp("#{HOSTSFILE}","#{HOSTSFILE_TMP}.old")  # Fairly-safe atomic replacement
  FileUtils.cp("#{HOSTSFILE_TMP}.new","#{HOSTSFILE}")
else
  pp mappings  # if not called with 'dump', just show the mappings that would be added
end