Merge branch 'develop' of github.com:fritzy/SleekXMPP into develop

1 files changed, 65 insertions, 213 deletions

diff --git a/sleekxmpp/clientxmpp.py b/sleekxmpp/clientxmpp.py
index fb5b2087..ad127726 100644
--- a/sleekxmpp/clientxmpp.py
+++ b/sleekxmpp/clientxmpp.py
@@ -15,12 +15,14 @@ import hashlib
 import random
 import threading
 
+import sleekxmpp
 from sleekxmpp import plugins
 from sleekxmpp import stanza
+from sleekxmpp import features
 from sleekxmpp.basexmpp import BaseXMPP
-from sleekxmpp.stanza import Message, Presence, Iq
+from sleekxmpp.stanza import *
 from sleekxmpp.xmlstream import XMLStream, RestartStream
-from sleekxmpp.xmlstream import StanzaBase, ET
+from sleekxmpp.xmlstream import StanzaBase, ET, register_stanza_plugin
 from sleekxmpp.xmlstream.matcher import *
 from sleekxmpp.xmlstream.handler import *
 
@@ -38,9 +40,12 @@ log = logging.getLogger(__name__)
 class ClientXMPP(BaseXMPP):
 
     """
-    SleekXMPP's client class.
+    SleekXMPP's client class. ( Use only for good, not for evil.)
 
-    Use only for good, not for evil.
+    Typical Use:
+    xmpp = ClientXMPP('user@server.tld/resource', 'password')
+    xmpp.process(block=False) // when block is True, it blocks the current
+    //                           thread. False by default.
 
     Attributes:
 
@@ -81,15 +86,19 @@ class ClientXMPP(BaseXMPP):
                 "xmlns='%s'" % self.default_ns)
         self.stream_footer = "</stream:stream>"
 
-        self.features = []
-        self.registered_features = []
+        self.features = set()
+        self._stream_feature_handlers = {}
+        self._stream_feature_order = []
 
         #TODO: Use stream state here
         self.authenticated = False
         self.sessionstarted = False
         self.bound = False
         self.bindfail = False
-        self.add_event_handler('connected', self.handle_connected)
+
+        self.add_event_handler('connected', self._handle_connected)
+
+        self.register_stanza(StreamFeatures)
 
         self.register_handler(
                 Callback('Stream Features',
@@ -102,32 +111,11 @@ class ClientXMPP(BaseXMPP):
                              'jabber:iq:roster')),
                          self._handle_roster))
 
-        self.register_feature(
-            "<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls' />",
-            self._handle_starttls, True)
-        self.register_feature(
-            "<mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl' />",
-            self._handle_sasl_auth, True)
-        self.register_feature(
-            "<bind xmlns='urn:ietf:params:xml:ns:xmpp-bind' />",
-            self._handle_bind_resource)
-        self.register_feature(
-            "<session xmlns='urn:ietf:params:xml:ns:xmpp-session' />",
-            self._handle_start_session)
-
-    def handle_connected(self, event=None):
-        #TODO: Use stream state here
-        self.authenticated = False
-        self.sessionstarted = False
-        self.bound = False
-        self.bindfail = False
-        self.schedule("session timeout checker", 15,
-                      self._session_timeout_check)
-
-    def _session_timeout_check(self):
-        if not self.session_started_event.isSet():
-            log.debug("Session start has taken more than 15 seconds")
-            self.disconnect(reconnect=self.auto_reconnect)
+        # Setup default stream features
+        self.register_plugin('feature_starttls')
+        self.register_plugin('feature_mechanisms')
+        self.register_plugin('feature_bind')
+        self.register_plugin('feature_session')
 
     def connect(self, address=tuple(), reattempt=True, use_tls=True):
         """
@@ -168,18 +156,23 @@ class ClientXMPP(BaseXMPP):
 
                     addresses = {}
                     intmax = 0
+                    topprio = 65535
+                    for answer in answers:
+                        topprio = min(topprio, answer.priority)
                     for answer in answers:
-                        intmax += answer.priority
-                        addresses[intmax] = (answer.target.to_text()[:-1],
+                        if answer.priority == topprio:
+                            intmax += answer.weight
+                            addresses[intmax] = (answer.target.to_text()[:-1],
                                              answer.port)
+
                     #python3 returns a generator for dictionary keys
-                    priorities = [x for x in addresses.keys()]
-                    priorities.sort()
+                    items = [x for x in addresses.keys()]
+                    items.sort()
 
                     picked = random.randint(0, intmax)
-                    for priority in priorities:
-                        if picked <= priority:
-                            address = addresses[priority]
+                    for item in items:
+                        if picked <= item:
+                            address = addresses[item]
                             break
 
         if not address:
@@ -189,19 +182,22 @@ class ClientXMPP(BaseXMPP):
         return XMLStream.connect(self, address[0], address[1],
                                  use_tls=use_tls, reattempt=reattempt)
 
-    def register_feature(self, mask, pointer, breaker=False):
+    def register_feature(self, name, handler, restart=False, order=5000):
         """
         Register a stream feature.
 
         Arguments:
-            mask    -- An XML string matching the feature's element.
-            pointer -- The function to execute if the feature is received.
-            breaker -- Indicates if feature processing should halt with
+            name    -- The name of the stream feature.
+            handler -- The function to execute if the feature is received.
+            restart -- Indicates if feature processing should halt with
                        this feature. Defaults to False.
+            order   -- The relative ordering in which the feature should
+                       be negotiated. Lower values will be attempted
+                       earlier when available.
         """
-        self.registered_features.append((MatchXMLMask(mask),
-                                         pointer,
-                                         breaker))
+        self._stream_feature_handlers[name] = (handler, restart)
+        self._stream_feature_order.append((order, name))
+        self._stream_feature_order.sort()
 
     def update_roster(self, jid, name=None, subscription=None, groups=[],
                             block=True, timeout=None, callback=None):
@@ -273,179 +269,35 @@ class ClientXMPP(BaseXMPP):
         else:
             return self._handle_roster(response, request=True)
 
-    def _handle_stream_features(self, features):
-        """
-        Process the received stream features.
-
-        Arguments:
-            features -- The features stanza.
-        """
-        # Record all of the features.
-        self.features = []
-        for sub in features.xml:
-            self.features.append(sub.tag)
-
-        # Process the features.
-        for sub in features.xml:
-            for feature in self.registered_features:
-                mask, handler, halt = feature
-                if mask.match(sub):
-                    if handler(sub) and halt:
-                        # Don't continue if the feature was
-                        # marked as a breaker.
-                        return True
-
-    def _handle_starttls(self, xml):
-        """
-        Handle notification that the server supports TLS.
-
-        Arguments:
-            xml -- The STARTLS proceed element.
-        """
-        if not self.use_tls:
-            return False
-        elif not self.authenticated and self.ssl_support:
-            tls_ns = 'urn:ietf:params:xml:ns:xmpp-tls'
-            self.add_handler("<proceed xmlns='%s' />" % tls_ns,
-                             self._handle_tls_start,
-                             name='TLS Proceed',
-                             instream=True)
-            self.send_xml(xml, now=True)
-            return True
-        else:
-            log.warning("The module tlslite is required to log in" +\
-                            " to some servers, and has not been found.")
-            return False
-
-    def _handle_tls_start(self, xml):
-        """
-        Handle encrypting the stream using TLS.
-
-        Restarts the stream.
-        """
-        log.debug("Starting TLS")
-        if self.start_tls():
-            raise RestartStream()
-
-    def _handle_sasl_auth(self, xml):
-        """
-        Handle authenticating using SASL.
-
-        Arguments:
-            xml -- The SASL mechanisms stanza.
-        """
-        if self.use_tls and \
-           '{urn:ietf:params:xml:ns:xmpp-tls}starttls' in self.features:
-            return False
-
-        log.debug("Starting SASL Auth")
-        sasl_ns = 'urn:ietf:params:xml:ns:xmpp-sasl'
-        self.add_handler("<success xmlns='%s' />" % sasl_ns,
-                         self._handle_auth_success,
-                         name='SASL Sucess',
-                         instream=True)
-        self.add_handler("<failure xmlns='%s' />" % sasl_ns,
-                         self._handle_auth_fail,
-                         name='SASL Failure',
-                         instream=True)
-
-        sasl_mechs = xml.findall('{%s}mechanism' % sasl_ns)
-        if sasl_mechs:
-            for sasl_mech in sasl_mechs:
-                self.features.append("sasl:%s" % sasl_mech.text)
-            if 'sasl:PLAIN' in self.features and self.boundjid.user:
-                if sys.version_info < (3, 0):
-                    user = bytes(self.boundjid.user)
-                    password = bytes(self.password)
-                else:
-                    user = bytes(self.boundjid.user, 'utf-8')
-                    password = bytes(self.password, 'utf-8')
-
-                auth = base64.b64encode(b'\x00' + user + \
-                                        b'\x00' + password).decode('utf-8')
-
-                self.send("<auth xmlns='%s' mechanism='PLAIN'>%s</auth>" % (
-                    sasl_ns,
-                    auth),
-                    now=True)
-            elif 'sasl:ANONYMOUS' in self.features and not self.boundjid.user:
-                self.send("<auth xmlns='%s' mechanism='%s' />" % (
-                    sasl_ns,
-                    'ANONYMOUS'),
-                    now=True)
-            else:
-                log.error("No appropriate login method.")
-                self.disconnect()
-        return True
-
-    def _handle_auth_success(self, xml):
-        """
-        SASL authentication succeeded. Restart the stream.
-
-        Arguments:
-            xml -- The SASL authentication success element.
-        """
-        self.authenticated = True
-        self.features = []
-        raise RestartStream()
-
-    def _handle_auth_fail(self, xml):
-        """
-        SASL authentication failed. Disconnect and shutdown.
+    def _handle_connected(self, event=None):
+        #TODO: Use stream state here
+        self.authenticated = False
+        self.sessionstarted = False
+        self.bound = False
+        self.bindfail = False
+        self.features = set()
 
-        Arguments:
-            xml -- The SASL authentication failure element.
-        """
-        log.info("Authentication failed.")
-        self.event("failed_auth", direct=True)
-        self.disconnect()
+        def session_timeout():
+            if not self.session_started_event.isSet():
+                log.debug("Session start has taken more than 15 seconds")
+                self.disconnect(reconnect=self.auto_reconnect)
 
-    def _handle_bind_resource(self, xml):
-        """
-        Handle requesting a specific resource.
+        self.schedule("session timeout checker", 15, session_timeout)
 
-        Arguments:
-            xml -- The bind feature element.
-        """
-        log.debug("Requesting resource: %s" % self.boundjid.resource)
-        xml.clear()
-        iq = self.Iq(stype='set')
-        if self.boundjid.resource:
-            res = ET.Element('resource')
-            res.text = self.boundjid.resource
-            xml.append(res)
-        iq.append(xml)
-        response = iq.send(now=True)
-
-        bind_ns = 'urn:ietf:params:xml:ns:xmpp-bind'
-        self.set_jid(response.xml.find('{%s}bind/{%s}jid' % (bind_ns,
-                                                             bind_ns)).text)
-        self.bound = True
-        log.info("Node set to: %s" % self.boundjid.full)
-        session_ns = 'urn:ietf:params:xml:ns:xmpp-session'
-        if "{%s}session" % session_ns not in self.features or self.bindfail:
-            log.debug("Established Session")
-            self.sessionstarted = True
-            self.session_started_event.set()
-            self.event("session_start")
-
-    def _handle_start_session(self, xml):
+    def _handle_stream_features(self, features):
         """
-        Handle the start of the session.
+        Process the received stream features.
 
         Arguments:
-            xml -- The session feature element.
+            features -- The features stanza.
         """
-        if self.authenticated and self.bound:
-            iq = self.makeIqSet(xml)
-            response = iq.send(now=True)
-            log.debug("Established Session")
-            self.sessionstarted = True
-            self.session_started_event.set()
-            self.event("session_start")
-        else:
-            # Bind probably hasn't happened yet.
-            self.bindfail = True
+        for order, name in self._stream_feature_order:
+            if name in features['features']:
+                handler, restart = self._stream_feature_handlers[name]
+                if handler(features) and restart:
+                    # Don't continue if the feature requires
+                    # restarting the XML stream.
+                    return True
 
     def _handle_roster(self, iq, request=False):
         """