#!/usr/bin/env python

from datetime import datetime

import argparse
import subprocess
import time
import json
import sys
import shlex



# Generic function to run a command, removes top item from output as this is always either 'listing vhosts', 'listing queue' etc
def get_command(command):

    command = shlex.split(command)
    
    proc = subprocess.Popen(command,stdout=subprocess.PIPE,stderr=subprocess.PIPE,shell=False)
    output = proc.communicate()
    
    # Output[0] comes as a tuple so items are first split at \n before first and last items are removed because first items descripes the output and last item is and empty item
    out = output[0].split('\n')[1:-1]
    
    return out



# Returns a list of rabbitmq virtual hosts
def get_vhosts():

    command = '/usr/sbin/rabbitmqctl list_vhosts'
    vhosts = get_command(command)

    return vhosts



# Generates autodiscovery json output for virtualhost and its queues
def get_queues():

    vhosts = get_vhosts()

    json_list = []

    for vhost in vhosts:
        
        command = '/usr/sbin/rabbitmqctl list_queues -p ' + vhost + ' name'
        queues = get_command(command)
        
        # only if queues are not empty
        if queues:  

            for queue in queues:
                
                inst = {}
                inst = {"{#VHOST}":vhost, "{#QUEUE}":queue}
                json_list.append(inst)

    result = {"data":json_list}
    print json.dumps(result,sort_keys=True,indent=2)
                


# Generates autodiscovery json output for virtualhost and its queues
def get_exchanges():

    vhosts = get_vhosts()

    json_list = []
    
    for vhost in vhosts:

        command = '/usr/sbin/rabbitmqctl list_exchanges -p ' + vhost + ' name'

        exchanges = get_command(command)

        # default and empty exchanges should be ignored 
        default_and_empty_exchanges = ['amq.direct', 'amq.fanout', 'amq.match', 'amq.topic', 'amq.headers', 'amq.rabbitmq.trace', 'amq.rabbitmq.log', '']
        
        # only if exchanges are not empty
        if exchanges:
    
            for exchange in exchanges:
                
                # Ensures that non of the exchanges from default_and_empty_exchanges are included
                if all(x != exchange for x in default_and_empty_exchanges):

                    inst = {}
                    inst = {"{#VHOST}":vhost, "{#EXCHANGE}":exchange}
                    json_list.append(inst)

    result = {"data":json_list}
    print json.dumps(result,sort_keys=True,indent=2)



if __name__ == '__main__':

    try:

        parser= argparse.ArgumentParser()

        parser.add_argument('-lh', '--list_vhosts', action='store_true', help='lists virtual hosts')
        parser.add_argument('-lq', '--list_queues', action='store_true', help='lists virtual hosts and their queues')
        parser.add_argument('-le', '--list_exchanges', action='store_true', help='lists virtual hosts and their exchanges')

        args = parser.parse_args()

        if args.list_vhosts:
            vhosts = get_vhosts()
            for vhost in vhosts:
                print vhost

        elif args.list_queues:
            get_queues()

        elif args.list_exchanges:
            get_exchanges()

        else:
            parser.print_help()
            sys.exit(1)

    except Exception as e:
        print e