#!/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, can have a prefix like mybucket/prefix", :type => String, :multi => true, :required => true
  opt :max_objects, "Max number of objects to list before stopping to save bandwidth", :default => -1


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

    options[:buckets].each do |url|
      split = url.split('/')
      bucket = split[0]
      prefix = ""
      if (split[1])
        prefix = url[(split[0].length+1)..-1]
      end
      count = 0
      connection.directories.get(bucket, prefix: prefix).files.map do |file|
        count = count +1
        if (options[:max_objects]>0 && count>options[:max_objects])
          break
        end
      end
      if (options[:max_objects]>0 && count>options[:max_objects])
        event = event(url, "objectCount", count, "count was bigger than threshold #{options[:max_objects]}", "warning")
        report(event)
      else
        event = event(url, "objectCount", count, "All objects counted, threshold=#{options[:max_objects]}", "ok")
        report(event)
      end
    end
  end

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

Riemann::Tools::S3Metrics.run