diff options
| -rw-r--r-- | sleekxmpp/clientxmpp.py | 15 | ||||
| -rw-r--r-- | sleekxmpp/features/feature_mechanisms/mechanisms.py | 19 | ||||
| -rw-r--r-- | sleekxmpp/features/feature_starttls/starttls.py | 20 | ||||
| -rw-r--r-- | sleekxmpp/test/sleektest.py | 1 | ||||
| -rw-r--r-- | sleekxmpp/xmlstream/xmlstream.py | 1 |
5 files changed, 43 insertions, 13 deletions
diff --git a/sleekxmpp/clientxmpp.py b/sleekxmpp/clientxmpp.py index 69e7db6c..3313328a 100644 --- a/sleekxmpp/clientxmpp.py +++ b/sleekxmpp/clientxmpp.py @@ -98,14 +98,12 @@ class ClientXMPP(BaseXMPP): self.add_event_handler('connected', self._handle_connected) self.add_event_handler('session_bind', self._handle_session_bind) + self.add_event_handler('stream_start', self._handle_stream_start) + self.add_event_handler('session_start', self._handle_session_start) self.register_stanza(StreamFeatures) self.register_handler( - Callback('Stream Features', - MatchXPath('{%s}features' % self.stream_ns), - self._handle_stream_features)) - self.register_handler( Callback('Roster Update', MatchXPath('{%s}iq/{%s}query' % ( self.default_ns, @@ -119,6 +117,15 @@ class ClientXMPP(BaseXMPP): self.register_plugin('feature_mechanisms', pconfig={'use_mech': sasl_mech} if sasl_mech else None) + def _handle_stream_start(self, root): + self.register_handler( + Callback('Stream Features', + MatchXPath('{%s}features' % self.stream_ns), + self._handle_stream_features)) + + def _handle_session_start(self, e): + self.remove_handler('Stream Features') + def connect(self, address=tuple(), reattempt=True, use_tls=True, use_ssl=False): """Connect to the XMPP server. diff --git a/sleekxmpp/features/feature_mechanisms/mechanisms.py b/sleekxmpp/features/feature_mechanisms/mechanisms.py index deff5d30..86dbb361 100644 --- a/sleekxmpp/features/feature_mechanisms/mechanisms.py +++ b/sleekxmpp/features/feature_mechanisms/mechanisms.py @@ -63,6 +63,17 @@ class feature_mechanisms(base_plugin): self.xmpp.register_stanza(stanza.Challenge) self.xmpp.register_stanza(stanza.Response) + self.xmpp.add_event_handler('stream_start', + self._handle_stream_start) + self.xmpp.add_event_handler('session_start', + self._handle_session_start) + + self.xmpp.register_feature('mechanisms', + self._handle_sasl_auth, + restart=True, + order=self.config.get('order', 100)) + + def _handle_stream_start(self, root): self.xmpp.register_handler( Callback('SASL Success', MatchXPath(stanza.Success.tag_name()), @@ -80,10 +91,10 @@ class feature_mechanisms(base_plugin): MatchXPath(stanza.Challenge.tag_name()), self._handle_challenge)) - self.xmpp.register_feature('mechanisms', - self._handle_sasl_auth, - restart=True, - order=self.config.get('order', 100)) + def _handle_session_start(self, e): + self.xmpp.remove_handler('SASL Success') + self.xmpp.remove_handler('SASL Failure') + self.xmpp.remove_handler('SASL Challenge') def _handle_sasl_auth(self, features): """ diff --git a/sleekxmpp/features/feature_starttls/starttls.py b/sleekxmpp/features/feature_starttls/starttls.py index 4e2b6621..a7f048b1 100644 --- a/sleekxmpp/features/feature_starttls/starttls.py +++ b/sleekxmpp/features/feature_starttls/starttls.py @@ -27,20 +27,30 @@ class feature_starttls(base_plugin): self.description = "STARTTLS Stream Feature" self.stanza = stanza - self.xmpp.register_handler( - Callback('STARTTLS Proceed', - MatchXPath(stanza.Proceed.tag_name()), - self._handle_starttls_proceed, - instream=True)) self.xmpp.register_feature('starttls', self._handle_starttls, restart=True, order=self.config.get('order', 0)) + self.xmpp.add_event_handler('stream_start', + self._handle_stream_start) + self.xmpp.add_event_handler('session_start', + self._handle_session_start) + self.xmpp.register_stanza(stanza.Proceed) self.xmpp.register_stanza(stanza.Failure) register_stanza_plugin(StreamFeatures, stanza.STARTTLS) + def _handle_stream_start(self, root): + self.xmpp.register_handler( + Callback('STARTTLS Proceed', + MatchXPath(stanza.Proceed.tag_name()), + self._handle_starttls_proceed, + instream=True)) + + def _handle_session_start(self, e): + self.xmpp.remove_handler('STARTTLS Proceed') + def _handle_starttls(self, features): """ Handle notification that the server supports TLS. diff --git a/sleekxmpp/test/sleektest.py b/sleekxmpp/test/sleektest.py index dd3df29a..3694c642 100644 --- a/sleekxmpp/test/sleektest.py +++ b/sleekxmpp/test/sleektest.py @@ -367,6 +367,7 @@ class SleekTest(unittest.TestCase): if skip: if socket != 'live': # Mark send queue as usable + self.xmpp.event('session_start') self.xmpp.session_started_event.set() # Clear startup stanzas self.xmpp.socket.next_sent(timeout=1) diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py index 4c8696b3..701ff240 100644 --- a/sleekxmpp/xmlstream/xmlstream.py +++ b/sleekxmpp/xmlstream/xmlstream.py @@ -1223,6 +1223,7 @@ class XMLStream(object): # as handshakes. self.stream_end_event.clear() self.start_stream_handler(root) + self.event('stream_start', root, direct=True) depth += 1 if event == b'end': depth -= 1 |