#!/usr/bin/env ruby
require 'rubygems'
require 'fog'
require 'date'

require File.expand_path('../../lib/riemann/tools', __FILE__)

$0 = __FILE__  # Let's not expose our AWS keys in the process list

class Riemann::Tools::AWS
  include Riemann::Tools
  
  opt :access_key, "AWS access key", :type => String
  opt :secret_key, "Secret access key", :type => String
  opt :region, "AWS region", :type => String, :default => 'eu-west-1'

  opt :retirement_critical, "Number of days before retirement. Defaults to 2", :default => 2
  opt :event_warning, "Number of days before event. Defaults to nil (i.e. when the event appears)", :default => nil

  def initialize
    @compute = Fog::Compute.new(:aws_access_key_id => opts[:access_key],
                                :aws_secret_access_key => opts[:secret_key],
                                :region => opts[:region],
                                :provider => 'AWS')
  end

  def tick 
    instance_status = @compute.describe_instance_status.body["instanceStatusSet"]
    status = instance_status.inject({}) do |acc,i|
      acc[i.delete("instanceId")] = i
      acc 
    end

    hosts = @compute.servers.select { |s| s.state == "running" }.
      inject([status, {}]) do |(status, acc), host| 
        acc[host.private_dns_name] = status.delete(host.id); [status, acc] 
      end[1]

    hosts.each do |host, status|
      status['eventsSet'].each do |event|
        before, after = ['notBefore', 'notAfter'].map { |k| Date.parse event[k].to_s if event[k] }

        ev = {:host => host,
              :service => "aws_instance_status",
              :description => "#{event['code']}\n\nstart #{event['notBefore']}\nend #{event['notAfter']}\n\n#{event['description']}",
              :state => "ok",
              :ttl => 300}

        ev2 = if (event['code'] == 'instance-retirement') and
          Date.today >= before-opts[:retirement_critical]
          {:state => "critical"}
        elsif opts[:event_warning] and Date.today >= before-opts[:event_warning]
          {:state => "warning"}
        else
          {:state => "warning"}
        end

        report ev.merge(ev2)
      end
    end
  end
end
  
Riemann::Tools::AWS.run