#! /usr/bin/env ruby
#
# check-sqs-messages
#
# DESCRIPTION:
#   This plugin checks the number of messages in an Amazon Web Services SQS queue.
#
# OUTPUT:
#   plain-text
#
# PLATFORMS:
#   Linux
#
# DEPENDENCIES:
#   gem: aws-sdk-v2
#   gem: sensu-plugin
#
# USAGE:
#   check-sqs-messages -q my_queue -w 500 -c 1000
#   check-sqs-messages -p queue_prefix_ -W 100 -C 50
#
# NOTES:
#
# LICENSE:
#   Copyright (c) 2013, Justin Lambert <jlambert@letsevenup.com>
#   Released under the same terms as Sensu (the MIT license); see LICENSE
#   for details.
#

require 'sensu-plugin/check/cli'
require 'sensu-plugins-aws'
require 'aws-sdk'

#
# Check SQS Messages
#
class SQSMsgs < Sensu::Plugin::Check::CLI
  include Common

  option :aws_region,
         short: '-r AWS_REGION',
         long: '--aws-region REGION',
         description: 'AWS Region (defaults to us-east-1).',
         default: 'us-east-1'

  option :queue,
         short: '-q SQS_QUEUE',
         long: '--queue SQS_QUEUE',
         description: 'A comma seperated list of the SQS queue(s) you want to check the number of messages for',
         default: ''

  option :prefix,
         short: '-p PREFIX',
         long: '--prefix PREFIX',
         description: 'The prefix of the queues you want to check the number of messages for',
         default: ''

  option :metric,
         short: '-m METRIC',
         long: '--metric METRIC',
         description: 'The metric of the queues you want to check the number of messages for',
         default: 'ApproximateNumberOfMessages'

  option :warn_over,
         short: '-w WARN_OVER',
         long: '--warnnum WARN_OVER',
         description: 'Number of messages in the queue considered to be a warning',
         default: -1,
         proc: proc(&:to_i)

  option :crit_over,
         short: '-c CRIT_OVER',
         long: '--critnum CRIT_OVER',
         description: 'Number of messages in the queue considered to be critical',
         default: -1,
         proc: proc(&:to_i)

  option :warn_under,
         short: '-W WARN_UNDER',
         long: '--warnunder WARN_UNDER',
         description: 'Minimum number of messages in the queue considered to be a warning',
         default: -1,
         proc: proc(&:to_i)

  option :crit_under,
         short: '-C CRIT_UNDER',
         long: '--critunder CRIT_UNDER',
         description: 'Minimum number of messages in the queue considered to be critical',
         default: -1,
         proc: proc(&:to_i)

  def aws_config
    { access_key_id: config[:aws_access_key],
      secret_access_key: config[:aws_secret_access_key],
      region: config[:aws_region] }
  end

  def run
    Aws.config.update(aws_config)
    sqs = Aws::SQS::Resource.new

    if config[:prefix].empty?
      if config[:queue].empty?
        critical 'Error, either QUEUE or PREFIX must be specified'
      end

      warnings = []
      crits = []
      passing = []
      queues = config[:queue].split(',')
      queues.each do |q|
        url = sqs.get_queue_by_name(queue_name: q).url
        messages = sqs.client.get_queue_attributes(queue_url: url, attribute_names: ['All']).attributes[config[:metric]].to_i

        if (config[:crit_under] >= 0 && messages < config[:crit_under]) || (config[:crit_over] >= 0 && messages > config[:crit_over])
          crits << "#{messages} message(s) in #{q}"
        elsif (config[:warn_under] >= 0 && messages < config[:warn_under]) || (config[:warn_over] >= 0 && messages > config[:warn_over])
          warnings << "#{messages} message(s) in #{q}"
        else
          passing << "#{messages} message(s) in #{q}"
        end
      end
      if crits.any?
        critical crits.join(', ').to_s
      elsif warnings.any?
        warning warnings.join(', ').to_s
      else
        ok "all queue(s): #{queues} are OK"
      end
    else
      warn = false
      crit = false
      queues = []

      sqs.queues(queue_name_prefix: config[:prefix]).each do |q|
        messages = sqs.client.get_queue_attributes(queue_url: q.url, attribute_names: ['All']).attributes[config[:metric]].to_i
        queue_name = q.attributes['QueueArn'].split(':').last

        if (config[:crit_under] >= 0 && messages < config[:crit_under]) || (config[:crit_over] >= 0 && messages > config[:crit_over])
          crit = true
          queues << "#{messages} message(s) in #{queue_name} queue"
        elsif (config[:warn_under] >= 0 && messages < config[:warn_under]) || (config[:warn_over] >= 0 && messages > config[:warn_over])
          warn = true
          queues << "#{messages} message(s) in #{queue_name} queue"
        end
      end

      if crit
        critical queues.to_s
      elsif warn
        warning queues.to_s
      else
        ok "All queues matching prefix '#{config[:prefix]}' OK"
      end
    end
  end
end