#!/usr/bin/ruby

# what: Rudy EC2 startup script
#  who: delano@solutious.com
# when: 2009-02-25 (rev: 4)

# NOTE: This is a prototype version of this script. A cleaner version 
# with better documentation is forthcoming.

# Runs when an ec2 instance startups up.
# Grabs configuration from the run time user data and stores it locally.
# See Rudy::CLI::Machines#start_instance for the expected yaml format.

# Put this in /etc/init.d. Then:
#   * chmod 755 rudy-ec2-startup
#   * cd /etc/rc3.d
#   * ln -s ../init.d/rudy-ec2-startup S17rudy
#   * cd /etc/rc4.d                    
#   * ln -s ../init.d/rudy-ec2-startup S17rudy
#   * cd /etc/rc5.d                    
#   * ln -s ../init.d/rudy-ec2-startup S17rudy

require 'yaml'
require 'resolv'

LOGFILE = '/var/log/rudy-ec2-startup'
USERDATA = 'http://169.254.169.254/2008-02-01/user-data'
METADATA = 'http://169.254.169.254/2008-02-01/meta-data'

#LOGFILE = '/tmp/rudy-ec2-startup'
#USERDATA = 'http://127.0.0.1/2008-02-01/user-data'
#METADATA = 'http://127.0.0.1/2008-02-01/meta-data'

METADATAPARAMS = ['instance-id', 'instance-type']

Rudy::CONFIG_DIR=File.join('.rudy')
Rudy::CONFIG_FILE=File.join('config')

module Rudy
  module EC2Startup
    extend self
    
    VERSION = 4.freeze
    
    def process_config
      begin
        File.unlink(LOGFILE) if File.exists?(LOGFILE)
        log "RUDY (rev#{VERSION}) is in the house"
        
        log "Grabing configuration..."
        config_yaml = run("curl -s #{USERDATA}")
  
        log "Grabing meta-data..."
        metadata = get_metadata
        
        log "Processing metadata..."
        METADATAPARAMS.each do |name|
          default = name.gsub('instance-', '')
          value = (metadata[name] && !metadata[name].empty?) ? metadata[name] : default
          log "  ---> #{name}: #{value}"
          write_to_file("/etc/ec2/#{name}", value, "w")
        end
        
        log "Processing userdata..."
        
        if !config_yaml || config_yaml.empty?
          raise "Nothing to process (did you supply user-data when you created the instance)" 
        end
        
        config = YAML::load(config_yaml)
        config ||= {}
        
        # Print to STDOUT and NOT to the log. 
        # We don't want sensitive config values available in the log
        puts "CONFIG: " << config_yaml
        
        zone = config[:zone] ? config[:zone].to_s.downcase : "zone"
        environment = config[:environment] ? config[:environment].to_s.downcase : "env"
        role = config[:role] ? config[:role].to_s.downcase : "role"
        position = config[:position] ? config[:position].to_s.downcase : "00"
        
        hostname = [zone, environment, role, position].join('-')
        log "Setting hostname: #{hostname}"
        `hostname #{hostname}`
        
        config[:hosts] ||= {}
        config[:hosts][hostname] = '127.0.0.1'
        
        if config[:userdata] && config[:userdata].is_a?(Hash)
          # TODO: How to we get the path to any user's home directory? 
          # (when we're not running as that user.) 
          config[:userdata].each_pair do |n,hash|
            fileparts = (n == :root) ? ['/root'] : [user_home_dir(n)]
            fileparts << Rudy::CONFIG_DIR
            dirpath = File.join(fileparts)
            filepath = File.join(dirpath, Rudy::CONFIG_FILE)
            
            # Backwards compatability
            if !File.directory?(dirpath)
              log "Deleting the deprecated #{dirpath} config"
              File.unlink(dirpath) 
            end
            
            unless File.exists?(dirpath)
              log "Creating #{dirpath}"
              Dir.mkdir(dirpath, 0700)
            end
            
            log "Writing to #{filepath}"
            user_data = {
              :userdata => hash
            }
            write_to_file(filepath, user_data.to_yaml, "w")
          end
        end
        
        if config[:hosts] && config[:hosts].is_a?(Hash)
          unless read_file('/etc/hosts') =~ /----- RUDY SAYS -----/
            log "Updating /etc/hosts..."
            write_to_file("/etc/hosts", "\n# ----- RUDY SAYS -----\n", :append)
            config[:hosts].each_pair do |name, address|
              # Respect existing entries
              next if read_file('/etc/hosts') =~ /#{name}/
      
              ip_address = (address !~ /^\d.+/) ? Resolv.getaddress(address) : address
      
              log "  ---> #{name}: #{ip_address}"
              write_to_file("/etc/hosts", "\n#{ip_address}\t#{name}\n", :append)
            end
          end
        end
        
        log "Done!"
  
      rescue => ex
        log "ERROR: #{ex.message}"
      end
    end
    
    
    def get_metadata
      metadata = {}

      begin
        METADATAPARAMS.each do |param|
          log "  ---> #{param}"
          metadata[param.to_s] = run("curl -s #{METADATA}/#{param}")
          metadata[param.to_s] = nil if metadata[param.to_s].empty?
        end
      rescue => ex
        log("Problem getting metadata: #{ex.message}")
      end
      metadata
    end

    def run(command, input='')
      IO.popen(command, 'r+') do |io|
        io.read
      end
    end  

    def get_formatted_time
      t = Time.now
      t_out = t.strftime("%H:%M:%S%p (%m/%d/%Y)")
    end

    def write_to_file(filename, s, type)
      type = (type == :append) ? 'a' : 'w'
      f = File.open(filename,type)
      f.puts s
      f.close
    end

    def read_file(path)
      read_file_to_array(path).join('')
    end

    def read_file_to_array(path)
      contents = []
      return contents unless File.exists?(path)
      open(path, 'r') do |l|
        contents = l.readlines
      end
      contents
    end
    
    def user_home_dir(user)
      (`su #{user} -c "echo \\$HOME"` || '').chomp
    end
    
    def log(s)
      msg = "#{get_formatted_time}: #{s}"
      write_to_file(LOGFILE, msg, :append)
      puts msg
    end
    
  end
end

Rudy::EC2Startup.process_config