#!/usr/bin/env ruby
# frozen_string_literal: true

Process.setproctitle($PROGRAM_NAME)

require 'riemann/tools'

$0 = __FILE__

module Riemann
  module Tools
    class 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
          connection = if options[:aws_access] && options[:aws_secret]
                         Fog::Storage.new({
                                            provider: 'AWS',
                                            aws_access_key_id: options[:aws_access],
                                            aws_secret_access_key: options[:aws_secret],
                                            region: options[:aws_region],
                                          })
                       else
                         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 = ''
          prefix = url[(split[0].length + 1)..] if split[1]
          count = 0
          connection.directories.get(bucket, prefix: prefix).files.map do |_file|
            count += 1
            break if options[:max_objects].positive? && count > options[:max_objects]
          end
          event = if options[:max_objects].positive? && count > options[:max_objects]
                    event(
                      url, 'objectCount', count, "count was bigger than threshold #{options[:max_objects]}",
                      'warning',
                    )
                  else
                    event(url, 'objectCount', count, "All objects counted, threshold=#{options[:max_objects]}", 'ok')
                  end
          report(event)
        end
      end

      private

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

Riemann::Tools::S3Metrics.run