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

require 'riemann/tools'

$0 = __FILE__

class Riemann::Tools::S3Metrics
  include Riemann::Tools
  require 'fog'
  require 'time'

  opt :fog_credentials_file, "Fog credentials file", :type => String
  opt :fog_credential, "Fog credentials to use", :type => String
  opt :aws_access, "AWS Access Key", :type => String
  opt :aws_secret, "AWS Secret Key", :type => String
  opt :aws_region, "AWS Region", :type => String, :default => "eu-west-1"
  opt :buckets, "Buckets to pull metrics from, multi=true", :type => String, :multi => true, :required => true
  opt :statistic, "Statistic to retrieve, multi=true, e.g. --statistic=Average --statistic=Maximum", :type => String, :multi => true, :required => true


  def base_metrics
    # get last 60 seconds
    start_time = (Time.now.utc - 3600 * 24 * 1).iso8601
    end_time = Time.now.utc.iso8601

    # The base query that all metrics would get
    metric_base = {
        "Namespace" => "AWS/S3",
        "StartTime" => start_time,
        "EndTime" => end_time,
        "Period" => 3600,
        "MetricName" => "NumberOfObjects",
    }

    metric_base
  end


  def tick
    if options[:fog_credentials_file]
      Fog.credentials_path = options[:fog_credentials_file]
      Fog.credential = options[:fog_credential].to_sym
      connection = Fog::AWS::CloudWatch.new
    else
      if options[:aws_access] && options[:aws_secret]
        connection = Fog::AWS::CloudWatch.new({
                                                  :aws_access_key_id => options[:aws_access],
                                                  :aws_secret_access_key => options[:aws_secret],
                                                  :region => options[:aws_region]
                                              })
      else
        connection = Fog::AWS::CloudWatch.new({
                                                  :use_iam_profile => true,
                                                  :region => options[:aws_region]
                                              })
      end
    end

    options[:statistic].each do |statistic|
      options[:buckets].each do |bucket|

        metric_base_options = base_metrics
        metric_base_options["Statistics"] = statistic
        metric_base_options["Dimensions"] = [
            {"Name" => "BucketName", "Value" => bucket},
            {"Name" => "StorageType", "Value" => "AllStorageTypes"}]

        result = connection.get_metric_statistics(metric_base_options)
        if result.body["GetMetricStatisticsResult"]["Datapoints"].empty?
          next
        end
          result.body["GetMetricStatisticsResult"]["Datapoints"][0].keys.sort.each do |stat_type|
          next if stat_type == "Unit"
          next if stat_type == "Timestamp"

          unit = result.body["GetMetricStatisticsResult"]["Datapoints"][0]["Unit"]
          metric = result.body["GetMetricStatisticsResult"]["Datapoints"][0][stat_type]
          event = event(bucket, result.body["GetMetricStatisticsResult"]["Label"], stat_type, unit, metric)
          report(event)
        end
      end
    end
  end

  private
  def event(bucket, label, metric_type, stat_type, metric, unit=nil)
    event = {
        host: "bucket_#{bucket}",
        service: "s3.#{label}.#{metric_type}.#{stat_type}",
        ttl: 300,
        description: "#{bucket} #{metric_type} #{stat_type} (#{unit})",
        tags: ["s3_metrics"],
        metric: metric
    }
  end
end

Riemann::Tools::S3Metrics.run