#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

from optparse import OptionParser, OptionGroup, IndentedHelpFormatter
import sys
import socket
import os
import locale
from qmf.console import Session, BrokerURL

_verbose   = False
_quiet     = False
_durable   = False
_dellink   = False
_srclocal  = False
_transport = "tcp"
_ack       = 0
_connTimeout = 10

class RouteManager:
    def __init__(self, localBroker):
        self.brokerList = {}
        self.local = BrokerURL(localBroker)
        self.remote  = None
        self.qmf = Session()
        self.broker = self.qmf.addBroker(localBroker, _connTimeout)
        self.broker._waitForStable()
        self.agent = self.broker.getBrokerAgent()

    def disconnect(self):
        try:
            if self.broker:
                self.qmf.delBroker(self.broker)
                self.broker = None
            while len(self.brokerList):
                b = self.brokerList.popitem()
                if b[0] != self.local.name():
                    self.qmf.delBroker(b[1])
        except:
            pass  # ignore errors while shutting down

    def getLink(self):
        links = self.agent.getObjects(_class="link")
        for link in links:
            if self.remote.match(link.host, link.port):
                return link
        return None

    def addLink(self, remoteBroker, mech="PLAIN"):
        self.remote = BrokerURL(remoteBroker)
        if self.local.match(self.remote.host, self.remote.port):
            raise Exception("Linking broker to itself is not permitted")

        brokers = self.agent.getObjects(_class="broker")
        broker = brokers[0]
        link = self.getLink()
        if link == None:
            res = broker.connect(self.remote.host, self.remote.port, _durable,
                                 mech, self.remote.authName or "", self.remote.authPass or "",
                                 _transport)
            if _verbose:
                print "Connect method returned:", res.status, res.text

    def delLink(self, remoteBroker):
        self.remote = BrokerURL(remoteBroker)
        brokers = self.agent.getObjects(_class="broker")
        broker = brokers[0]
        link = self.getLink()
        if link == None:
            raise Exception("Link not found")

        res = link.close()
        if _verbose:
            print "Close method returned:", res.status, res.text

    def listLinks(self):
        links = self.agent.getObjects(_class="link")
        if len(links) == 0:
            print "No Links Found"
        else:
            print
            print "Host            Port    Transport Durable  State             Last Error"
            print "============================================================================="
            for link in links:
                print "%-16s%-8d%-13s%c     %-18s%s" % \
                (link.host, link.port, link.transport, YN(link.durable), link.state, link.lastError)

    def mapRoutes(self):
        print
        print "Finding Linked Brokers:"

        self.brokerList[self.local.name()] = self.broker
        print "    %s... Ok" % self.local

        added = True
        while added:
            added = False
            links = self.qmf.getObjects(_class="link")
            for link in links:
                url = BrokerURL("%s:%d" % (link.host, link.port))
                if url.name() not in self.brokerList:
                    print "    %s..." % url.name(),
                    try:
                        b = self.qmf.addBroker("%s:%d" % (link.host, link.port), _connTimeout)
                        self.brokerList[url.name()] = b
                        added = True
                        print "Ok"
                    except Exception, e:
                        print e

        print
        print "Dynamic Routes:"
        bridges = self.qmf.getObjects(_class="bridge", dynamic=True)
        fedExchanges = []
        for bridge in bridges:
            if bridge.src not in fedExchanges:
                fedExchanges.append(bridge.src)
        if len(fedExchanges) == 0:
            print "  none found"
        print

        for ex in fedExchanges:
            print "  Exchange %s:" % ex
            pairs = []
            for bridge in bridges:
                if bridge.src == ex:
                    link = bridge._linkRef_
                    fromUrl = "%s:%s" % (link.host, link.port)
                    toUrl = bridge.getBroker().getUrl()
                    found = False
                    for pair in pairs:
                        if pair.matches(fromUrl, toUrl):
                            found = True
                    if not found:
                        pairs.append(RoutePair(fromUrl, toUrl))
            for pair in pairs:
                print "    %s" % pair
            print

        print "Static Routes:"
        bridges = self.qmf.getObjects(_class="bridge", dynamic=False)
        if len(bridges) == 0:
            print "  none found"
        print

        for bridge in bridges:
            link = bridge._linkRef_
            fromUrl = "%s:%s" % (link.host, link.port)
            toUrl = bridge.getBroker().getUrl()
            leftType = "ex"
            rightType = "ex"
            if bridge.srcIsLocal:
                arrow = "=>"
                left = bridge.src
                right = bridge.dest
                if bridge.srcIsQueue:
                    leftType = "queue"
            else:
                arrow = "<="
                left = bridge.dest
                right = bridge.src
                if bridge.srcIsQueue:
                    rightType = "queue"

            if bridge.srcIsQueue:
                print "  %s(%s=%s) %s %s(%s=%s)" % \
                    (toUrl, leftType, left, arrow, fromUrl, rightType, right)
            else:
                print "  %s(%s=%s) %s %s(%s=%s) key=%s" % \
                    (toUrl, leftType, left, arrow, fromUrl, rightType, right, bridge.key)
        print

        while len(self.brokerList):
            b = self.brokerList.popitem()
            if b[0] != self.local.name():
                self.qmf.delBroker(b[1])

    def addRoute(self, remoteBroker, exchange, routingKey, tag, excludes, mech="PLAIN", dynamic=False):
        if dynamic and _srclocal:
            raise Exception("--src-local is not permitted on dynamic routes")

        self.addLink(remoteBroker, mech)
        link = self.getLink()
        if link == None:
            raise Exception("Link failed to create")

        bridges = self.agent.getObjects(_class="bridge")
        for bridge in bridges:
            if bridge.linkRef == link.getObjectId() and \
                    bridge.dest == exchange and bridge.key == routingKey and not bridge.srcIsQueue:
                if not _quiet:
                    raise Exception("Duplicate Route - ignoring: %s(%s)" % (exchange, routingKey))
                sys.exit(0)

        if _verbose:
            print "Creating inter-broker binding..."
        res = link.bridge(_durable, exchange, exchange, routingKey, tag, excludes, False, _srclocal, dynamic, _ack)
        if res.status != 0:
            raise Exception(res.text)
        if _verbose:
            print "Bridge method returned:", res.status, res.text

    def addQueueRoute(self, remoteBroker, exchange, queue):
        self.addLink(remoteBroker)
        link = self.getLink()
        if link == None:
            raise Exception("Link failed to create")

        bridges = self.agent.getObjects(_class="bridge")
        for bridge in bridges:
            if bridge.linkRef == link.getObjectId() and \
                    bridge.dest == exchange and bridge.src == queue and bridge.srcIsQueue:
                if not _quiet:
                    raise Exception("Duplicate Route - ignoring: %s(%s)" % (exchange, queue))
                sys.exit(0)

        if _verbose:
            print "Creating inter-broker binding..."
        res = link.bridge(_durable, queue, exchange, "", "", "", True, _srclocal, False, _ack)
        if res.status != 0:
            raise Exception(res.text)
        if _verbose:
            print "Bridge method returned:", res.status, res.text

    def delQueueRoute(self, remoteBroker, exchange, queue):
        self.remote = BrokerURL(remoteBroker)
        link = self.getLink()
        if link == None:
            if not _quiet:
                raise Exception("No link found from %s to %s" % (self.remote.name(), self.local.name()))
            sys.exit(0)

        bridges = self.agent.getObjects(_class="bridge")
        for bridge in bridges:
            if bridge.linkRef == link.getObjectId() and \
                    bridge.dest == exchange and bridge.src == queue and bridge.srcIsQueue:
                if _verbose:
                    print "Closing bridge..."
                res = bridge.close()
                if res.status != 0:
                    raise Exception("Error closing bridge: %d - %s" % (res.status, res.text))
                if len(bridges) == 1 and _dellink:
                    link = self.getLink()
                    if link == None:
                        sys.exit(0)
                    if _verbose:
                        print "Last bridge on link, closing link..."
                    res = link.close()
                    if res.status != 0:
                        raise Exception("Error closing link: %d - %s" % (res.status, res.text))
                sys.exit(0)
        if not _quiet:
            raise Exception("Route not found")

    def delRoute(self, remoteBroker, exchange, routingKey, dynamic=False):
        self.remote = BrokerURL(remoteBroker)
        link = self.getLink()
        if link == None:
            if not _quiet:
                raise Exception("No link found from %s to %s" % (self.remote.name(), self.local.name()))
            sys.exit(0)

        bridges = self.agent.getObjects(_class="bridge")
        for bridge in bridges:
            if bridge.linkRef == link.getObjectId() and bridge.dest == exchange and bridge.key == routingKey \
                    and bridge.dynamic == dynamic:
                if _verbose:
                    print "Closing bridge..."
                res = bridge.close()
                if res.status != 0:
                    raise Exception("Error closing bridge: %d - %s" % (res.status, res.text))
                if len(bridges) == 1 and _dellink:
                    link = self.getLink()
                    if link == None:
                        sys.exit(0)
                    if _verbose:
                        print "Last bridge on link, closing link..."
                    res = link.close()
                    if res.status != 0:
                        raise Exception("Error closing link: %d - %s" % (res.status, res.text))
                return
        if not _quiet:
            raise Exception("Route not found")

    def listRoutes(self):
        links   = self.qmf.getObjects(_class="link")
        bridges = self.qmf.getObjects(_class="bridge")

        for bridge in bridges:
            myLink = None
            for link in links:
                if bridge.linkRef == link.getObjectId():
                    myLink = link
                    break
            if myLink != None:
                if bridge.dynamic:
                    keyText = "<dynamic>"
                else:
                    keyText = bridge.key
                print "%s %s:%d %s %s" % (self.local.name(), myLink.host, myLink.port, bridge.dest, keyText)

    def clearAllRoutes(self):
        links   = self.qmf.getObjects(_class="link")
        bridges = self.qmf.getObjects(_class="bridge")

        for bridge in bridges:
            if _verbose:
                myLink = None
                for link in links:
                    if bridge.linkRef == link.getObjectId():
                        myLink = link
                        break
                if myLink != None:
                    print "Deleting Bridge: %s:%d %s %s... " % (myLink.host, myLink.port, bridge.dest, bridge.key),
            res = bridge.close()
            if res.status != 0:
                print "Error: %d - %s" % (res.status, res.text)
            elif _verbose:
                print "Ok"

        if _dellink:
            links = self.qmf.getObjects(_class="link")
            for link in links:
                if _verbose:
                    print "Deleting Link: %s:%d... " % (link.host, link.port),
                res = link.close()
                if res.status != 0:
                    print "Error: %d - %s" % (res.status, res.text)
                elif _verbose:
                    print "Ok"

class RoutePair:
    def __init__(self, fromUrl, toUrl):
        self.fromUrl = fromUrl
        self.toUrl = toUrl
        self.bidir = False

    def __repr__(self):
        if self.bidir:
            delimit = "<=>"
        else:
            delimit = " =>"
        return "%s %s %s" % (self.fromUrl, delimit, self.toUrl)

    def matches(self, fromUrl, toUrl):
        if fromUrl == self.fromUrl and toUrl == self.toUrl:
            return True
        if toUrl == self.fromUrl and fromUrl == self.toUrl:
            self.bidir = True
            return True
        return False
        

def YN(val):
    if val == 1:
        return 'Y'
    return 'N'

##
## Main Program
##


class JHelpFormatter(IndentedHelpFormatter):
    """Format usage and description without stripping newlines from usage strings
    """

    def format_usage(self, usage):
        return usage


    def format_description(self, description):
        if description:
            return description + "\n"
        else:
            return ""

def usage(parser):
    parser.print_help()
    exit(-1)

usage = """
Usage:  qpid-route [OPTIONS] dynamic add <dest-broker> <src-broker> <exchange> [tag] [exclude-list]
        qpid-route [OPTIONS] dynamic del <dest-broker> <src-broker> <exchange>

        qpid-route [OPTIONS] route add   <dest-broker> <src-broker> <exchange> <routing-key> [tag] [exclude-list] [mechanism]
        qpid-route [OPTIONS] route del   <dest-broker> <src-broker> <exchange> <routing-key>
        qpid-route [OPTIONS] queue add   <dest-broker> <src-broker> <exchange> <queue>
        qpid-route [OPTIONS] queue del   <dest-broker> <src-broker> <exchange> <queue>
        qpid-route [OPTIONS] route list  [<dest-broker>]
        qpid-route [OPTIONS] route flush [<dest-broker>]
        qpid-route [OPTIONS] route map   [<broker>]

        qpid-route [OPTIONS] link add  <dest-broker> <src-broker>
        qpid-route [OPTIONS] link del  <dest-broker> <src-broker>
        qpid-route [OPTIONS] link list [<dest-broker>]"""

description = """
ADDRESS syntax:   

      [username/password@] hostname
      ip-address [:<port>]"""
         
parser = OptionParser(usage=usage,
                      description=description,
                      formatter=JHelpFormatter())

parser.add_option("--timeout", action="store", type="int", default="10", metavar="SECS", help="Maximum time to wait for broker connection (in seconds)")
parser.add_option("-v", "--verbose", action="store_true", help="Verbose output")
parser.add_option("-q", "--quiet", action="store_true", help="Quiet output, don't print duplicate warnings")
parser.add_option("-d", "--durable", action="store_true", help="Added configuration shall be durable")

parser.add_option("-e", "--del-empty-link", action="store_true", help="Delete link after deleting last route on the link")
parser.add_option("-s", "--src-local", action="store_true", help="Make connection to source broker (push route)")

parser.add_option("--ack", action="store", type="int", metavar="N", help="Acknowledge transfers over the bridge in batches of N")
parser.add_option("-t", "--transport", action="store", type="string", default="tcp", metavar="<transport>", help="Transport to use for links, defaults to tcp")

opts, encArgs = parser.parse_args()

try:
    encoding = locale.getpreferredencoding()
    args = [a.decode(encoding) for a in encArgs]
except:
    args = encArgs

if opts.timeout:
    _connTimeout = opts.timeout
    if _connTimeout == 0:
        _connTimeout = None
    
if opts.verbose:
    _verbose = True

if opts.quiet:
    _quiet = True

if opts.durable:
    _durable = True

if opts.del_empty_link:
    _dellink = True

if opts.src_local:
    _srclocal = true

if opts.transport:
    _transport = opts.transport

if opts.ack:
    _ack = opts.ack

nargs = len(args)
if nargs < 2:
    usage(parser)

if nargs == 2:
    localBroker = socket.gethostname()
else:
    if _srclocal:
        localBroker = args[3]
        remoteBroker = args[2]
    else:
        localBroker = args[2]
        if nargs > 3:
            remoteBroker = args[3]

group = args[0]
cmd   = args[1]

rm = None
try:
    rm = RouteManager(localBroker)
    if group == "link":
        if cmd == "add":
            if nargs != 4:
                usage(parser)
            rm.addLink(remoteBroker)
        elif cmd == "del":
            if nargs != 4:
                usage(parser)
            rm.delLink(remoteBroker)
        elif cmd == "list":
            rm.listLinks()

    elif group == "dynamic":
        if cmd == "add":
            if nargs < 5 or nargs > 7:
                usage(parser)

            tag = ""
            excludes = ""
            mech = "PLAIN"
            if nargs > 5: tag = args[5]     
            if nargs > 6: excludes = args[6]     
            rm.addRoute(remoteBroker, args[4], "", tag, excludes, mech, dynamic=True)
        elif cmd == "del":
            if nargs != 5:
                usage(parser)
            else:
                rm.delRoute(remoteBroker, args[4], "", dynamic=True)

    elif group == "route":
        if cmd == "add":
            if nargs < 6 or nargs > 9:
                usage(parser)

            tag = ""
            excludes = ""
            mech = "PLAIN"
            if nargs > 6: tag = args[6]     
            if nargs > 7: excludes = args[7]
            if nargs > 8: mech = args[8]
            rm.addRoute(remoteBroker, args[4], args[5], tag, excludes, mech, dynamic=False)
        elif cmd == "del":
            if nargs != 6:
                usage(parser)
            rm.delRoute(remoteBroker, args[4], args[5], dynamic=False)
        elif cmd == "map":
            rm.mapRoutes()
        else:
            if cmd == "list":
                rm.listRoutes()
            elif cmd == "flush":
                rm.clearAllRoutes()
            else:
                usage(parser)

    elif group == "queue":
        if nargs != 6:
            usage(parser)
        if cmd == "add":
            rm.addQueueRoute(remoteBroker, exchange=args[4], queue=args[5])
        elif cmd == "del":
            rm.delQueueRoute(remoteBroker, exchange=args[4], queue=args[5])
        else:
            usage(parser)

except Exception,e:
    if rm:
        rm.disconnect()  # try to release broker resources
    print "Failed: %s - %s" % (e.__class__.__name__, e)
    sys.exit(1)

rm.disconnect()
