#!/usr/bin/env python
from optparse import OptionParser
import shlex
import subprocess
import sys


class RabbitCmdWrapper(object):
    """So basically this just runs rabbitmqctl commands and returns parsed output.
       Typically this means you need root privs for this to work.

       Made this it's own class so it could be used in other monitoring tools
       if desired."""

    @classmethod
    def list_connections(cls):
        args = shlex.split("sudo /usr/sbin/rabbitmqctl list_connections")
        cmd_result = subprocess.check_output(args).strip()
        results = cls._parse_list_results(cmd_result)
        return results

    @classmethod
    def list_queues(cls):
        args = shlex.split('sudo /usr/sbin/rabbitmqctl list_queues')
        cmd_result = subprocess.check_output(args).strip()
        results = cls._parse_list_results(cmd_result)
        return results

    @classmethod
    def status(cls):
        args = shlex.split('sudo /usr/sbin/rabbitmqctl status')
        cmd_result = subprocess.check_output(args).strip()
        results = cls._parse_list_results(cmd_result)
        return results

    @classmethod
    def _parse_list_results(cls, result_string):
        results = result_string.strip().split('\n')
        #remove text fluff
        if "Listing connections ..." in results: results.remove("Listing connections ...")
        if "Listing queues ..." in results: results.remove("Listing queues ...")
        return_data = []
        for row in results:
            return_data.append(row.split('\t'))
        return return_data


def check_connection_count(critical=0, warning=0):
    """Checks to make sure the numbers of connections are within parameters."""
    try:
        count = len(RabbitCmdWrapper.list_connections())
        if count >= critical:
            print "CRITICAL - Connection Count %d" % count
            sys.exit(2)
        elif count >= warning:
            print "WARNING - Connection Count %d" % count
            sys.exit(1)
        else:
            print "OK - Connection Count %d" % count
    except Exception, err:
        print "CRITICAL - %s" % err


def check_queues_count(critical=1000, warning=1000):
    """
    A blanket check to make sure all queues are within count parameters.
    TODO: Possibly break this out so test can be done on individual queues.
    """
    try:
        critical_q = []
        warning_q = []
        ok_q = []
        results = RabbitCmdWrapper.list_queues()
        for queue in results:
            count = int(queue[1])
            if count >= critical:
                critical_q.append("%s: %s" % (queue[0], count))
            elif count >= warning:
                warning_q.append("%s: %s" % (queue[0], count))
            else:
                ok_q.append("%s: %s" % (queue[0], count))
        if critical_q:
            print "CRITICAL - %s" % ", ".join(critical_q)
            sys.exit(2)
        elif warning_q:
            print "WARNING - %s" % ", ".join(warning_q)
            sys.exit(1)
        else:
            print "OK - %s" % ", ".join(ok_q)
            sys.exit(0)
    except Exception, err:
        print "CRITICAL - %s" % err
        sys.exit(2)

def check_mem_usage(critical=75, warning=50):
    """Check to make sure the RAM usage of rabbitmq process does not exceed 50%% of its max"""
    try:
        results = RabbitCmdWrapper.status()

        for idx,val in enumerate(results):
          if "memory," in str(val):
              mem_used_raw = str(results[idx + 1])
          if "vm_memory_limit" in str(val):
              mem_limit_raw = str(val)

        memory_used = float(filter(str.isdigit, mem_used_raw))
        memory_limit = float(filter(str.isdigit, mem_limit_raw))
        percent_usage = int(memory_used/memory_limit * 100)

        if percent_usage > critical:
            print "CRITICAL - RABBITMQ RAM USAGE at %s%% of max" % percent_usage
            sys.exit(2)
        elif percent_usage > warning:
            print "WARNING - RABBITMQ RAM USAGE at %s%% of max" % percent_usage
            sys.exit(1)
        else:
            print "OK - RABBITMQ RAM USAGE OK at %s%% of max" % percent_usage
            sys.exit(0)
    except Exception, err:
        print "Critical - %s" % err
        sys.exit(2)

USAGE = """Usage: ./check_rabbitmq -a [action] -C [critical] -W [warning]
           Actions:
           - connection_count
             checks the number of connection in rabbitmq's list_connections
           - queues_count
             checks the count in each of the queues in rabbitmq's list_queues
           - mem_usage
             checks to ensure mem usage of rabbitmq process does not exceed 50%"""

if __name__ == "__main__":
    parser = OptionParser(USAGE)
    parser.add_option("-a", "--action", dest="action",
                      help="Action to Check")
    parser.add_option("-C", "--critical", dest="critical",
                      type="int", help="Critical Threshold")
    parser.add_option("-W", "--warning", dest="warning",
                      type="int", help="Warning Threshold")
    (options, args) = parser.parse_args()

    if options.action == "connection_count":
        check_connection_count(options.critical, options.warning)
    elif options.action == "queues_count":
        check_queues_count(options.critical, options.warning)
    elif options.action == "mem_usage":
        check_mem_usage(options.critical, options.warning)
    else:
        print "Invalid action: %s" % options.action
        print USAGE