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


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

$0 = __FILE__

class Riemann::Tools::AWSBilling
  include Riemann::Tools

  opt :fog_credentials_file, "Fog credentials file", :type => String
  opt :fog_credential, "Fog credentials to use", :type => String

  opt :access_key, "AWS access key", :type => String
  opt :secret_key, "Secret access key", :type => String
  opt :services, "AWS services: AmazonEC2  AmazonS3  AWSDataTransfer", :type => :strings, :multi => true, :default => ["AmazonEC2", "AmazonS3", "AWSDataTransfer"]

  opt :time_start, "Start time in seconds of the metrics period (2hrs ago default)", :type => Integer, :default => 7200
  opt :time_end, "End time in seconds of the metrics period ", :type => Integer, :default => 60


  def initialize
    if options[:fog_credentials_file]
      Fog.credentials_path = opts[:fog_credentials_file]
      Fog.credential = opts[:fog_credential].to_sym
      @cloudwatch = Fog::AWS::CloudWatch.new
    else
      @cloudwatch = Fog::AWS::CloudWatch.new(:aws_secret_access_key => opts[:secret_key], :aws_access_key_id => opts[:access_key])
      @start_time = (Time.now.utc - opts[:time_start]).iso8601
      @end_time = (Time.now.utc - opts[:time_end]).iso8601
    end
  end

  def tick
    opts[:services].each do |service|
      data = @cloudwatch.get_metric_statistics({
                                                 'Statistics' => ["Maximum"],
                                                 'StartTime' => @start_time,
                                                 'EndTime' => @end_time,
                                                 'Period' => 3600,
                                                 'Unit' => "None",
                                                 'MetricName' => "EstimatedCharges",
                                                 'Namespace' => "AWS/Billing",
                                                 'Dimensions' => [
                                                                  {
                                                                    'Name' => "ServiceName",
                                                                    'Value' => service
                                                                  },
                                                                  {
                                                                    'Name' => "Currency",
                                                                    'Value' => "USD"
                                                                  }
                                                                 ]
                                               }).body['GetMetricStatisticsResult']['Datapoints']


      data.each do |metrics|
        name = "AWScloudwatch.Billing." + service
        value = metrics["Maximum"]
        timestamp = metrics["Timestamp"].to_i

        event = {
          host: nil,
          service: name,
          time: timestamp,
          description: "AWS Estimate Charges for #{service}",
          tags: ["aws_billing"],
          state: "ok",
          metric: value
        }

        report event
      end
    end
  end
end

Riemann::Tools::AWSBilling.run