diff sat/core/patches.py @ 2691:1ecceac3df96

plugin XEP-0198: Stream Management implementation: - hooks can now be set in stream onElement and send methods - xmllog refactored to use new hooks - client.isConnected now uses transport.connected method - fixed reconnection, SàT will now try to reconnect indefinitely until it success, unresolvable failure happen (e.g. invalid certificate), or explicit disconnection is requested (or a plugin change this behaviour) - new triggers: "stream_hooks", "disconnecting", "disconnected", and "xml_init" (replace "XML Initialized")
author Goffi <goffi@goffi.org>
date Sun, 18 Nov 2018 15:49:46 +0100
parents e9cd473a2f46
children 00d905e1b0ef
line wrap: on
line diff
--- a/sat/core/patches.py	Sat Nov 10 10:16:38 2018 +0100
+++ b/sat/core/patches.py	Sun Nov 18 15:49:46 2018 +0100
@@ -1,10 +1,19 @@
-from twisted.words.protocols.jabber import xmlstream
+from twisted.words.protocols.jabber import xmlstream, sasl, client as tclient
 from twisted.internet import ssl
 from wokkel import client
+from sat.core.constants import Const as C
+from sat.core.log import getLogger
 
-"""This module apply monkey patches to Twisted and Wokkel to handle certificate validation
-   during XMPP connection"""
+log = getLogger(__name__)
 
+"""This module apply monkey patches to Twisted and Wokkel
+   First part handle certificate validation during XMPP connectionand are temporary
+   (until merged upstream).
+   Second part add a trigger point to send and onElement method of XmlStream
+   """
+
+
+## certificate validation patches
 
 class TLSInitiatingInitializer(xmlstream.TLSInitiatingInitializer):
     check_certificate = True
@@ -21,13 +30,15 @@
 
 class XMPPClient(client.XMPPClient):
 
-    def __init__(self, jid, password, host=None, port=5222, check_certificate=True):
+    def __init__(self, jid, password, host=None, port=5222,
+                 check_certificate=True):
         self.jid = jid
         self.domain = jid.host.encode('idna')
         self.host = host
         self.port = port
 
-        factory = HybridClientFactory(jid, password, check_certificate)
+        factory = HybridClientFactory(
+            jid, password, check_certificate=check_certificate)
 
         client.StreamManager.__init__(self, factory)
 
@@ -39,6 +50,7 @@
 
 
 class HybridAuthenticator(client.HybridAuthenticator):
+    res_binding = True
 
     def __init__(self, jid, password, check_certificate):
         xmlstream.ConnectAuthenticator.__init__(self, jid.host)
@@ -53,11 +65,83 @@
         tlsInit.check_certificate = self.check_certificate
         xs.initializers = [client.client.CheckVersionInitializer(xs),
                            tlsInit,
-                           client.CheckAuthInitializer(xs)]
+                           CheckAuthInitializer(xs, self.res_binding)]
+
+
+# XmlStream triggers
+
+
+class XmlStream(xmlstream.XmlStream):
+    """XmlStream which allows to add hooks"""
+
+    def __init__(self, authenticator):
+        xmlstream.XmlStream.__init__(self, authenticator)
+        # hooks at this level should not modify content
+        # so it's not needed to handle priority as with triggers
+        self._onElementHooks = []
+        self._sendHooks = []
+
+    def addHook(self, hook_type, callback):
+        """Add a send or receive hook"""
+        conflict_msg = (u"Hook conflict: can't add {hook_type} hook {callback}"
+            .format(hook_type=hook_type, callback=callback))
+        if hook_type == C.STREAM_HOOK_RECEIVE:
+            if callback not in self._onElementHooks:
+                self._onElementHooks.append(callback)
+            else:
+                log.warning(conflict_msg)
+        elif hook_type == C.STREAM_HOOK_SEND:
+            if callback not in self._sendHooks:
+                self._sendHooks.append(callback)
+            else:
+                log.warning(conflict_msg)
+        else:
+            raise ValueError(u"Invalid hook type: {hook_type}"
+                .format(hook_type=hook_type))
+
+    def onElement(self, element):
+        for hook in self._onElementHooks:
+            hook(element)
+        xmlstream.XmlStream.onElement(self, element)
+
+    def send(self, obj):
+        for hook in self._sendHooks:
+            hook(obj)
+        xmlstream.XmlStream.send(self, obj)
+
+
+# Binding activation (needed for stream management, XEP-0198)
+
+
+class CheckAuthInitializer(client.CheckAuthInitializer):
+
+    def __init__(self, xs, res_binding):
+        super(CheckAuthInitializer, self).__init__(xs)
+        self.res_binding = res_binding
+
+    def initialize(self):
+        # XXX: modification of client.CheckAuthInitializer which has optional
+        #      resource binding, and which doesn't do deprecated
+        #      SessionInitializer
+        if (sasl.NS_XMPP_SASL, 'mechanisms') in self.xmlstream.features:
+            inits = [(sasl.SASLInitiatingInitializer, True)]
+            if self.res_binding:
+                inits.append((tclient.BindInitializer, True)),
+
+            for initClass, required in inits:
+                init = initClass(self.xmlstream)
+                init.required = required
+                self.xmlstream.initializers.append(init)
+        elif (tclient.NS_IQ_AUTH_FEATURE, 'auth') in self.xmlstream.features:
+            self.xmlstream.initializers.append(
+                    tclient.IQAuthInitializer(self.xmlstream))
+        else:
+            raise Exception("No available authentication method found")
 
 
 def apply():
+    # certificate validation
     xmlstream.TLSInitiatingInitializer = TLSInitiatingInitializer
     client.XMPPClient = XMPPClient
-    client.HybridClientFactory = HybridClientFactory
-    client.HybridAuthenticator = HybridAuthenticator
+    # XmlStream triggers
+    xmlstream.XmlStreamFactory.protocol = XmlStream