#!/usr/bin/env python

# example commands:
# dump mq201 > topology.json
# diff topology.json mq201
# check topology.json
# check mq201
# graph topology.json > topology.dot
# graph topology.json | dot -Tpng > topology.png
# sync topology.json mq201

import urlparse
import urllib
import requests
import json
import copy
import sys

DEFAULT_PORT=55672

def graph_style(style):
    """generate style string from dict"""
    return ''.join([
        '[',
        ','.join(['%s="%s"' % (k, v) for (k, v) in style.iteritems()]),
        ']'
        ])

def graph_node(prefix, name, style):
    """generate a subgraph node for an exchange or queue"""
    template =  \
    'subgraph "cluster_{prefix}_{name}" {{\n' \
    'label="{name}";\n' \
    'color=transparent;\n' \
    '"{prefix}_{name}" {style};\n' \
    '}}\n\n'
    return template.format(prefix=prefix, 
            name=name,
            style=graph_style(style))

def graph_bind_nodes(source_type, source, destination_type, destination, style):
    """generate graph edge for rabbit binding"""
    return \
            '"{source_type}_{source}" -> "{dest_type}_{dest}" {style};\n'.format(
                    source_type=source_type,
                    source=source,
                    dest_type=destination_type,
                    dest=destination,
                    style=graph_style(style)
                    )

def host_port(address, default_port):
    """split host and port. return (host, port)"""
    hostport = address.rsplit(':', 1)
    if len(hostport) == 1:
        return (hostport[0], int(default_port))
    return (hostport[0], int(hostport[1]))

def filter_keys(dicts, *keys):
    """generate a copy of the list of dicts with keys removed"""
    for d in dicts:
        d = copy.copy(d)
        for key in keys:
            d.pop(key, None)
        yield d

def clean_exchanges(exchanges):
    """remove non-config keys from exchanges in list"""
    return filter_keys(exchanges, 'message_stats', 'message_stats_out', 'message_stats_in')

def skip_internal_exchanges(exchanges):
    """generator that skips the internal rabbit exchanges"""
    for exchange in exchanges:
        if exchange['name'] == '' or \
                exchange['name'].startswith('amq.') \
                or exchange['internal'] == True:
            continue
        yield exchange

def skip_non_permanent(things):
    """generator that skips non-durable and auto-delete exchanges or queues"""
    # works for queues and exchanges
    for thing in things:
        if thing['durable'] == True and thing['auto_delete'] == False:
            yield thing

def skip_empty_source(bindings):
    """generator that skips bindings with no source"""
    for binding in bindings:
        if binding['source'] == '':
            continue
        yield binding

def clean_queues(queues):
    """remove non-config keys from queues in list"""
    return filter_keys(queues,
            'node',
            'consumers',
            'messages_details',
            'messages', 
            'backing_queue_status',
            'message_stats',
            'messages_unacknowledged_details',
            'memory',
            'messages_ready_details',
            'messages_unacknowledged',
            'messages_ready',
            'idle_since',
            'active_consumers',
            'policy',
            'slave_nodes',
            'status'
            )

def clean_bindings(bindings):
    """remove non-config keys from bindings in list"""
    return filter_keys(bindings, 'properties_key')

def find_unbound_queues(bindings, queues):
    """find queues with no bindings"""
    bound_queues = set()
    unbound_queues = set()
    for binding in bindings:
        if binding['destination_type'] == 'queue':
            bound_queues.add(binding['destination'])
    for queue in queues:
        if queue['name'] not in bound_queues:
            unbound_queues.add(queue['name'])
    return list(unbound_queues)

def find_unbound_exchanges(bindings, queues, exchanges):
    """find exchanges with no bindings or dead-letter queues"""
    bound_exchanges = set()
    unbound_exchanges = set()
    for binding in bindings:
        bound_exchanges.add(binding['source'])
        if binding['destination_type'] == 'exchange':
            bound_exchanges.add(binding['destination'])
    for queue in queues:
        if 'x-dead-letter-exchange' in queue['arguments']:
            bound_exchanges.add(queue['arguments']['x-dead-letter-exchange'])
    for exchange in exchanges:
        if exchange['name'] not in bound_exchanges:
            unbound_exchanges.add(exchange['name'])
    return list(unbound_exchanges)

def find_no_consumers_no_ttl(queues):
    """find queues with no consumers and no ttl"""
    no = []
    for queue in queues:
        if queue.get('consumers', None) == 0 and \
                queue['arguments'].get('x-message-ttl') == None:
            no.append(queue['name'])
    return no

def find_no_consumers_no_dlx(queues):
    """find queues with no consumers and no dead letter exchange"""
    no = []
    for queue in queues:
        if queue.get('consumers', None) == 0 and \
                queue['arguments'].get('x-dead-letter-exchange') == None:
            no.append(queue['name'])
    return no

def fetch_json(address, path, auth):
    """return decoded JSON from HTTP API"""
    (host, port) = host_port(address, DEFAULT_PORT)
    url = 'http://%s:%s%s' % (host, port, path)
    r = requests.get(url, auth=auth)
    return r.json()

def diff(expected, actual, get_key):
    """return dict with missing, extra, and different keys based on expected and actual lists"""
    exp_map = {}
    act_map = {}
    for ex in expected:
        exp_map[get_key(ex)] = ex
    for ac in actual:
        act_map[get_key(ac)] = ac
    exp_set = set(exp_map.keys())
    act_set = set(act_map.keys())
    missing = exp_set.difference(act_set)
    extra = act_set.difference(exp_set)
    different = []
    for key in exp_set.intersection(act_set):
        if exp_map[key] != act_map[key]:
            different.append(key)
    return dict(
            missing=list(missing),
            extra=list(extra),
            different=list(different)
            )

def get_remote_topology(address, auth):
    """get exchanges, bindings, and queues from rabbitmq server"""
    exchanges = fetch_json(address, '/api/exchanges', auth)
    exchanges = list(
            clean_exchanges(
                skip_internal_exchanges(
                    skip_non_permanent(exchanges)
                    )
                )
            )

    bindings = fetch_json(address, '/api/bindings', auth)
    bindings = list(
            clean_bindings(
                skip_empty_source(bindings)
                )
            )

    queues_raw = fetch_json(address, '/api/queues', auth)
    queues = list(
            skip_non_permanent(queues_raw)
            )

    topology = dict(
            exchanges=exchanges,
            bindings=bindings,
            queues=queues
            )
    return topology

def get_topology(file_or_host, auth):
    """try to load topology from local file or rabbitmq server"""
    try:
        with open(file_or_host) as f:
            return json.load(f)
    except IOError:
        pass
    return get_remote_topology(file_or_host, auth)

def probe(address, auth):
    """call overview api and follow redirect to determine port api is hosted on"""
    (host, port) = host_port(address, DEFAULT_PORT)
    url = 'http://%s:%s/api/overview' % (host, port)
    r = requests.get(url, auth=auth)
    url = urlparse.urlparse(r.url)
    return (url.hostname, url.port)

def put_api(host, port, auth, data, path, *args):
    """call api with  PUT method"""
    url = 'http://%s:%s' % (host, port)
    url = url + (path % tuple([urllib.quote(arg, '') for arg in args]))
    headers = {'Content-Type':'application/json'}
    data=json.dumps(data)
    return requests.put(url,
            data=data,
            auth=auth,
            headers=headers,
            )

def post_api(host, port, auth, data, path, *args):
    """call api with POST method"""
    url = 'http://%s:%s' % (host, port)
    url = url + (path % tuple([urllib.quote(arg, '') for arg in args]))
    headers = {'Content-Type':'application/json'}
    data=json.dumps(data)
    return requests.post(url,
            data=data,
            auth=auth,
            headers=headers,
            )

def create_exchange(host, port, auth, exchange):
    """use api to create an exchange"""
    return put_api(host, port, auth, exchange, '/api/exchanges/%s/%s',
            exchange['vhost'],
            exchange['name'])

def create_queue(host, port, auth, queue):
    """use api to create a queue"""
    return put_api(host, port, auth, queue, '/api/queues/%s/%s',
            queue['vhost'],
            queue['name'])

def create_binding(host, port, auth, binding):
    """use api to create a binding"""
    # TODO: support exchange->exchange binding?
    return post_api(host, port, auth, binding, '/api/bindings/%s/e/%s/q/%s',
            binding['vhost'],
            binding['source'],
            binding['destination'])

def command_check(args):
    auth = (args.user, args.password)
    if len(args.argv) == 1:
        topology = get_topology(args.argv[0], auth)
        bindings = topology['bindings']
        queues = topology['queues']
        exchanges = topology['exchanges']
        warnings = dict(
                unbound_queues=find_unbound_queues(bindings, queues),
                unbound_exchanges=find_unbound_exchanges(bindings, queues, exchanges),
                no_consumers_no_ttl=find_no_consumers_no_ttl(skip_non_permanent(queues)),
                no_consumers_no_dlx=find_no_consumers_no_dlx(skip_non_permanent(queues)),
                )
        return warnings

def command_dump(args):
    auth = (args.user, args.password)
    topology = get_remote_topology(args.argv[0], auth)
    topology['queues'] = list(clean_queues(topology['queues']))
    return topology

def command_diff(args):
    auth = (args.user, args.password)
    expected = get_topology(args.argv[0], auth)
    actual = get_topology(args.argv[1], auth)
    actual['queues'] = list(clean_queues(actual['queues']))

    def get_name(d):
        return d['name']
    def get_bind_name(d):
        return '%s->%s:%s' % (
                d['source'],
                d['destination'],
                d['destination_type']
                )

    warnings = {}
    ops = (
            ('exchanges', get_name),
            ('queues', get_name),
            ('bindings', get_bind_name),
            )
    for (thing, getter) in ops:
        warnings[thing] = diff(expected[thing], actual[thing], getter)
    return warnings

def command_sync(args):
    auth = (args.user, args.password)
    topology = get_topology(args.argv[0], auth)
    (host, port) = probe(args.argv[1], auth)
    address = '%s:%s' % (host, port)
    errors = []
    for exchange in topology['exchanges']:
        r = create_exchange(host, port, auth, exchange)
        if r.status_code != 204:
            errors.append((r.url, r.json()))
    for queue in topology['queues']:
        r = create_queue(host, port, auth, queue)
        if r.status_code != 204:
            errors.append((r.url, r.json()))
    for binding in topology['bindings']:
        r = create_binding(host, port, auth, binding)
        if r.status_code != 201:
            errors.append((r.url, r.json()))
    return errors

def command_graph(args):
    auth = (args.user, args.password)
    topology = get_topology(args.argv[0], auth)
    exchanges = [
            graph_node(
                prefix='X',
                name=exchange['name'],
                style=dict(label='', color='green', shape='ellipse'),
                )
            for exchange in topology['exchanges']
            ]
    queues = [
            graph_node(
                prefix='Q',
                name=queue['name'],
                style=dict(label='', color='red', shape='record'),
                )
            for queue in topology['queues']
            ]
    bindings = [
            graph_bind_nodes(
                'X',
                b['source'],
                'Q',
                b['destination'],
                style=dict(label=b['routing_key']),
                )
            for b in topology['bindings']
            ]
    dead_letter_bindings = []
    for queue in topology['queues']:
        if 'x-dead-letter-exchange' in queue['arguments']:
            dead_letter_bindings.append(dict(
                source=queue['name'],
                destination=queue['arguments']['x-dead-letter-exchange'],
                routing_key='')
                )
    bindings.extend([
            graph_bind_nodes(
                'Q',
                b['source'],
                'X',
                b['destination'],
                style=dict(label=b['routing_key'],style='dotted'),
                )
            for b in dead_letter_bindings
            ])
    return ''.join([
        'digraph {\n',
        'bgcolor=transparent;\n',
        'truecolor=true;\n',
        'rankdir=LR;\n',
        'node [style="filled"];\n\n', 
        ''.join(map(''.join, [exchanges, queues, bindings])),
        '}'])

def main(argv):
    import argparse
    parser = argparse.ArgumentParser(
            description='RabbitMQ Topology')
    parser.add_argument('command', choices=['check', 'diff', 'dump', 'graph', 'sync'])
    parser.add_argument('argv', nargs='*', help='arguments')
    parser.add_argument('-u', '--user', default='guest', help='basic auth username')
    parser.add_argument('-p', '--password', default='guest', help='basic auth password')
    args = parser.parse_args()
    command = None
    if args.command == 'check':
        command = command_check
    if args.command == 'dump':
        command = command_dump
    if args.command == 'diff':
        command = command_diff
    if args.command == 'sync':
        command = command_sync
    if args.command == 'graph':
        command = command_graph

    if command:
        result = command(args)
        if isinstance(result, basestring):
            sys.stdout.write(result)
        else:
            json.dump(result, sys.stdout, indent=2)
        print

if __name__ == '__main__':
    main(sys.argv)
