From eff3330e754cfbb496d6ba761d0a821375f575c8 Mon Sep 17 00:00:00 2001
From: Lance Stout <lancestout@gmail.com>
Date: Mon, 12 Dec 2011 22:17:07 -0800
Subject: [PATCH] Add support for incoming/outgoing filters.

A filter accepts and returns a stanza, but potentially modified.

To prevent sending/receiving a stanza, a filter may return None.

Incoming:
    self.add_filter('in', in_filter)

Outgoing:
    self.add_filter('out', out_filter)

Filters are applied in the order thay are added. However, you may
add an order parameter, which is the place in the list to insert the
filter:

    self.add_filter('in', in_filter, order=0)
---
 sleekxmpp/xmlstream/stanzabase.py |  2 +-
 sleekxmpp/xmlstream/xmlstream.py  | 29 ++++++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 389fe20..2f86430 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -1251,7 +1251,7 @@ class StanzaBase(ElementBase):
                          stanza sent immediately. Useful for stream
                          initialization. Defaults to ``False``.
         """
-        self.stream.send_raw(self.__str__(), now=now)
+        self.stream.send(self, now=now)
 
     def __copy__(self):
         """Return a copy of the stanza object that does not share the
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index 3e56908..b690103 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -35,7 +35,7 @@ except ImportError:
 import sleekxmpp
 from sleekxmpp.thirdparty.statemachine import StateMachine
 from sleekxmpp.xmlstream import Scheduler, tostring
-from sleekxmpp.xmlstream.stanzabase import StanzaBase, ET
+from sleekxmpp.xmlstream.stanzabase import StanzaBase, ET, ElementBase
 from sleekxmpp.xmlstream.handler import Waiter, XMLCallback
 from sleekxmpp.xmlstream.matcher import MatchXMLMask
 
@@ -268,6 +268,7 @@ class XMLStream(object):
         self.__handlers = []
         self.__event_handlers = {}
         self.__event_handlers_lock = threading.Lock()
+        self.__filters = {'in': [], 'out': []}
 
         self._id = 0
         self._id_lock = threading.Lock()
@@ -743,6 +744,18 @@ class XMLStream(object):
         """
         del self.__root_stanza[stanza_class]
 
+    def add_filter(self, mode, handler, order=None):
+        """Add a filter for incoming or outgoing stanzas.
+
+        These filters are applied before incoming stanzas are
+        passed to any handlers, and before outgoing stanzas
+        are put in the send queue.
+        """
+        if order:
+            self.__filters[mode].insert(order, handler)
+        else:
+            self.__filters[mode].append(handler)
+
     def add_handler(self, mask, pointer, name=None, disposable=False,
                     threaded=False, filter=False, instream=False):
         """A shortcut method for registering a handler using XML masks.
@@ -994,6 +1007,14 @@ class XMLStream(object):
             timeout = self.response_timeout
         if hasattr(mask, 'xml'):
             mask = mask.xml
+
+        if isinstance(data, ElementBase):
+            for filter in self.__filters['out']:
+                if data is not None:
+                    data = filter(data)
+            if data is None:
+                return
+
         data = str(data)
         if mask is not None:
             log.warning("Use of send mask waiters is deprecated.")
@@ -1254,6 +1275,12 @@ class XMLStream(object):
         # stanza type applies, a generic StanzaBase stanza will be used.
         stanza = self._build_stanza(xml)
 
+        for filter in self.__filters['in']:
+            if stanza is not None:
+                stanza = filter(stanza)
+        if stanza is None:
+            return
+
         # Match the stanza against registered handlers. Handlers marked
         # to run "in stream" will be executed immediately; the rest will
         # be queued.