#!/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.
#

import os
from optparse import OptionParser, OptionGroup
import sys
import locale
import socket
import re
from qpid.messaging import Connection

home = os.environ.get("QPID_TOOLS_HOME", os.path.normpath("/usr/share/qpid-tools"))
sys.path.append(os.path.join(home, "python"))

from qpidtoollibs import BrokerAgent
from qpidtoollibs import Display, Header, Sorter, YN, Commas, TimeLong


class Config:
    def __init__(self):
        self._host = "localhost"
        self._connTimeout = 10

config = Config()
conn_options = {}

def OptionsAndArguments(argv):
    """ Set global variables for options, return arguments """

    global config
    global conn_options

    usage = "%prog [options]"

    parser = OptionParser(usage=usage)

    parser.add_option("-b", "--broker",  action="store", type="string", default="localhost", metavar="<url>",
                      help="URL of the broker to query")
    parser.add_option("-t", "--timeout", action="store", type="int", default=10, metavar="<secs>",
                      help="Maximum time to wait for broker connection (in seconds)")
    parser.add_option("--sasl-mechanism", action="store", type="string", metavar="<mech>",
                      help="SASL mechanism for authentication (e.g. EXTERNAL, ANONYMOUS, PLAIN, CRAM-MD, DIGEST-MD5, GSSAPI). SASL automatically picks the most secure available mechanism - use this option to override.")
    parser.add_option("--ssl-certificate", action="store", type="string", metavar="<cert>", help="Client SSL certificate (PEM Format)")
    parser.add_option("--ssl-key", action="store", type="string", metavar="<key>", help="Client SSL private key (PEM Format)")
    parser.add_option("--ha-admin", action="store_true", help="Allow connection to a HA backup broker.")

    opts, args = parser.parse_args(args=argv)

    config._host = opts.broker
    config._connTimeout = opts.timeout

    if opts.sasl_mechanism:
        conn_options['sasl_mechanisms'] = opts.sasl_mechanism
    if opts.ssl_certificate:
        conn_options['ssl_certfile'] = opts.ssl_certificate
    if opts.ssl_key:
        conn_options['ssl_key'] = opts.ssl_key
    if opts.ha_admin:
        conn_options['client_properties'] = {'qpid.ha-admin' : 1}
    return args

class BrokerManager:
    def __init__(self):
        self.brokerName = None
        self.connection = None
        self.broker     = None
        self.cluster    = None

    def SetBroker(self, brokerUrl):
        self.url = brokerUrl
        self.connection = Connection.establish(self.url, **conn_options)
        self.broker = BrokerAgent(self.connection)

    def Disconnect(self):
        """ Release any allocated brokers.  Ignore any failures as the tool is
        shutting down.
        """
        try:
            connection.close()
        except:
            pass

    def Ping(self, args):
        for sequence in range(10):
            result = self.broker.echo(sequence, "ECHO BODY")
            if result['sequence'] != sequence:
                raise Exception("Invalid Sequence")


def main(argv=None):

    args = OptionsAndArguments(argv)
    bm   = BrokerManager()

    try:
        bm.SetBroker(config._host)
        bm.Ping(args)
        bm.Disconnect()
        return 0
    except KeyboardInterrupt:
        print
    except Exception,e:
        print "Failed: %s - %s" % (e.__class__.__name__, e)

    bm.Disconnect()   # try to deallocate brokers
    return 1

if __name__ == "__main__":
    sys.exit(main())