#!/usr/bin/env ruby
Process.setproctitle($0)

require 'riemann/tools'

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

class Riemann::Tools::AWS
  include Riemann::Tools
  require 'fog'
  require 'date'

  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
    if opts.has_key?('secret_key') and opts.has_key?('access_key')
      creds = {
        :aws_secret_access_key => opts[:secret_key],
        :aws_access_key_id => opts[:access_key]
      }
    else
      creds = { :use_iam_profile => true }
    end
    creds['region'] = opts[:region]
    creds['provider'] = 'AWS'
    @compute = Fog::Compute.new(creds)
  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