#!/usr/bin/env python_browser
#
# Example usage: ./disco.py items sleek@conference.jabber.org
#

import sys
import time
import logging
import getpass
import sleekxmpp
from optparse import OptionParser


class Disco(sleekxmpp.ClientXMPP):

    def __init__(self, jid, password, target, target_node='main', get=''):
        sleekxmpp.ClientXMPP.__init__(self, jid, password)
        self.add_event_handler("session_start", self.start)
        self.add_event_handler("disco_info", self.disco_info)
        self.add_event_handler("disco_items", self.disco_items)

        self.get = get
        self.target = target
        self.target_node = target_node
        self.results = {'identities': None,
                        'features': None,
                        'items': None}

        # Values for self.get to control which disco entities are reported
        self.info_types = ['', 'all', 'info', 'identities', 'features']
        self.identity_types = ['', 'all', 'info', 'identities']
        self.feature_types = ['', 'all', 'info', 'features']
        self.items_types = ['', 'all', 'items']

    def start(self, event):
        self.getRoster()
        self.sendPresence()
        if self.get in self.info_types:
            self['xep_0030'].getInfo(self.target, node=self.target_node)
        elif self.get in self.items_types:
            self['xep_0030'].getItems(self.target, node=self.target_node)
        else:
            logging.error("Invalid disco request type.")
            self.disconnect()

    def disco_info(self, iq):
        self.results['identities'] = iq['disco_info'].getIdentities()
        self.results['features'] = iq['disco_info'].getFeatures()
        if self.get in self.items_types:
            self['xep_0030'].getItems(self.target, node=self.target_node)
        else:
            self.print_results()

    def disco_items(self, iq):
        self.results['items'] = iq['disco_items'].getItems()
        self.print_results()

    def print_results(self):
        header = 'XMPP Service Discovery: %s' % self.target
        print header
        print '-' * len(header)
        if self.target_node != '':
            print 'Node: %s' % self.target_node
            print '-' * len(header)

        if self.get in self.identity_types:
            print 'Identities:'
            for identity in self.results['identities']:
                print '  - ', identity

        if self.get in self.feature_types:
            print 'Features:'
            for feature in self.results['features']:
                print '  - %s' % feature

        if self.get in self.items_types:
            print 'Items:'
            for item in self.results['items']:
                print '  - %s' % str(item)

        self.disconnect()

if __name__ == '__main__':
    optp = OptionParser()
    optp.usage = "Usage: %prog [options] info|items|identities|features <jid> [<node>]"
    optp.version = '%%prog 0.1'

    optp.add_option('-q','--quiet', help='set logging to ERROR',
                    action='store_const',
                    dest='loglevel',
                    const=logging.ERROR,
                    default=logging.ERROR)
    optp.add_option('-d','--debug', help='set logging to DEBUG',
                    action='store_const',
                    dest='loglevel',
                    const=logging.DEBUG,
                    default=logging.ERROR)
    optp.add_option('-v','--verbose', help='set logging to COMM',
                    action='store_const',
                    dest='loglevel',
                    const=5,
                    default=logging.ERROR)
    opts,args = optp.parse_args()

    logging.basicConfig(level=opts.loglevel, format='%(levelname)-8s %(message)s')

    if len(args) < 2:
        optp.print_help()
        exit()

    if len(args) == 2:
        args = (args[0], args[1], '')

    username = raw_input("Username: ")
    password = getpass.getpass("Password: ")

    xmpp = Disco(username, password, args[1], args[2], args[0])
    xmpp.registerPlugin('xep_0030')
    if xmpp.connect():
        xmpp.process(threaded=False)
    else:
        print("Unable to connect.")