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

2 files changed, 61 insertions, 30 deletions

diff --git a/sleekxmpp/plugins/xep_0066/oob.py b/sleekxmpp/plugins/xep_0066/oob.py
index 1f79492d..98cb81cd 100644
--- a/sleekxmpp/plugins/xep_0066/oob.py
+++ b/sleekxmpp/plugins/xep_0066/oob.py
@@ -9,6 +9,7 @@
 import logging
 
 from sleekxmpp.stanza import Message, Presence, Iq
+from sleekxmpp.exceptions import XMPPError
 from sleekxmpp.xmlstream import register_stanza_plugin
 from sleekxmpp.xmlstream.handler import Callback
 from sleekxmpp.xmlstream.matcher import StanzaPath
@@ -47,6 +48,9 @@ class xep_0066(base_plugin):
         self.description = 'Out-of-Band Transfer'
         self.stanza = stanza
 
+        self.url_handlers = {'global': self._default_handler,
+                             'jid': {}}
+
         register_stanza_plugin(Iq, stanza.OOBTransfer)
         register_stanza_plugin(Message, stanza.OOB)
         register_stanza_plugin(Presence, stanza.OOB)
@@ -62,6 +66,28 @@ class xep_0066(base_plugin):
         self.xmpp['xep_0030'].add_feature(stanza.OOBTransfer.namespace)
         self.xmpp['xep_0030'].add_feature(stanza.OOB.namespace)
 
+    def register_url_handler(self, jid=None, handler=None):
+        """
+        Register a handler to process download requests, either for all
+        JIDs or a single JID.
+
+        Arguments:
+            jid     -- If None, then set the handler as a global default.
+            handler -- If None, then remove the existing handler for the
+                       given JID, or reset the global handler if the JID
+                       is None.
+        """
+        if jid is None:
+            if handler is not None:
+                self.url_handlers['global'] = handler
+            else:
+                self.url_handlers['global'] = self._default_handler
+        else:
+            if handler is not None:
+                self.url_handlers['jid'][jid] = handler
+            else:
+                del self.url_handlers['jid'][jid]
+
     def send_oob(self, to, url, desc=None, ifrom=None, **iqargs):
         """
         Initiate a basic file transfer by sending the URL of
@@ -88,8 +114,41 @@ class xep_0066(base_plugin):
         iq['oob_transfer']['desc'] = desc
         return iq.send(**iqargs)
 
+    def _run_url_handler(self, iq):
+        """
+        Execute the appropriate handler for a transfer request.
+
+        Arguments:
+            iq -- The Iq stanza containing the OOB transfer request.
+        """
+        if iq['to'] in self.url_handlers['jid']:
+            return self.url_handlers['jid'][jid](iq)
+        else:
+            if self.url_handlers['global']:
+                self.url_handlers['global'](iq)
+            else:
+                raise XMPPError('service-unavailable')
+
+    def _default_handler(self, iq):
+        """
+        As a safe default, don't actually download files.
+
+        Register a new handler using self.register_url_handler to
+        screen requests and download files.
+
+        Arguments:
+            iq -- The Iq stanza containing the OOB transfer request.
+        """
+        raise XMPPError('service-unavailable')
+
     def _handle_transfer(self, iq):
-        """Handle receiving an out-of-band transfer request."""
+        """
+        Handle receiving an out-of-band transfer request.
+
+        Arguments:
+            iq -- An Iq stanza containing an OOB transfer request.
+        """
         log.debug('Received out-of-band data request for %s from %s:' % (
             iq['oob_transfer']['url'], iq['from']))
-        self.xmpp.event('oob_transfer', iq)
+        self._run_url_handler(iq)
+        iq.reply().send()
diff --git a/tests/test_stream_xep_0066.py b/tests/test_stream_xep_0066.py
index 3dbaf840..e3f2ddfa 100644
--- a/tests/test_stream_xep_0066.py
+++ b/tests/test_stream_xep_0066.py
@@ -40,33 +40,5 @@ class TestOOB(SleekTest):
 
         t.join()
 
-    def testReceiveOOB(self):
-        """Test receiving an OOB request."""
-        self.stream_start(plugins=['xep_0066', 'xep_0030'])
-
-        events = []
-
-        def receive_oob(iq):
-            events.append(iq['oob_transfer']['url'])
-
-        self.xmpp.add_event_handler('oob_transfer', receive_oob)
-
-        self.recv("""
-          <iq to="tester@localhost"
-              from="user@example.com"
-              type="set" id="1">
-            <query xmlns="jabber:iq:oob">
-              <url>http://github.com/fritzy/SleekXMPP/blob/master/README</url>
-              <desc>SleekXMPP README</desc>
-            </query>
-          </iq>
-        """)
-
-        time.sleep(0.1)
-
-        self.assertEqual(events,
-                ['http://github.com/fritzy/SleekXMPP/blob/master/README'],
-                'URL was not received: %s' % events)
-
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestOOB)