Fix handling forwarded stanzas to do proper lookups and deletions.

2 files changed, 19 insertions, 12 deletions

diff --git a/sleekxmpp/plugins/xep_0297/forwarded.py b/sleekxmpp/plugins/xep_0297/forwarded.py
index 7876967c..95703a2d 100644
--- a/sleekxmpp/plugins/xep_0297/forwarded.py
+++ b/sleekxmpp/plugins/xep_0297/forwarded.py
@@ -26,9 +26,14 @@ class XEP_0297(BasePlugin):
 
     def plugin_init(self):
         register_stanza_plugin(Message, Forwarded)
-        register_stanza_plugin(Forwarded, Message)
-        register_stanza_plugin(Forwarded, Presence)
-        register_stanza_plugin(Forwarded, Iq)
+
+        # While these are marked as iterable, that is just for
+        # making it easier to extract the forwarded stanza. There
+        # still can be only a single forwarded stanza.
+        register_stanza_plugin(Forwarded, Message, iterable=True)
+        register_stanza_plugin(Forwarded, Presence, iterable=True)
+        register_stanza_plugin(Forwarded, Iq, iterable=True)
+
         register_stanza_plugin(Forwarded, self.xmpp['xep_0203'].stanza.Delay)
 
         self.xmpp.register_handler(
diff --git a/sleekxmpp/plugins/xep_0297/stanza.py b/sleekxmpp/plugins/xep_0297/stanza.py
index 1cf02f74..8b97accc 100644
--- a/sleekxmpp/plugins/xep_0297/stanza.py
+++ b/sleekxmpp/plugins/xep_0297/stanza.py
@@ -6,6 +6,7 @@
     See the file LICENSE for copying permission.
 """
 
+from sleekxmpp.stanza import Message, Presence, Iq
 from sleekxmpp.xmlstream import ElementBase
 
 
@@ -16,12 +17,9 @@ class Forwarded(ElementBase):
     interfaces = set(['stanza'])
 
     def get_stanza(self):
-        if self.xml.find('{jabber:client}message') is not None:
-            return self['message']
-        elif self.xml.find('{jabber:client}presence') is not None:
-            return self['presence']
-        elif self.xml.find('{jabber:client}iq') is not None:
-            return self['iq']
+        for stanza in self:
+            if isinstance(stanza, (Message, Presence, Iq)):
+                return stanza
         return ''
 
     def set_stanza(self, value):
@@ -29,6 +27,10 @@ class Forwarded(ElementBase):
         self.append(value)
 
     def del_stanza(self):
-        del self['message']
-        del self['presence']
-        del self['iq']
+        found_stanzas = []
+        for stanza in self:
+            if isinstance(stanza, (Message, Presence, Iq)):
+                found_stanzas.append(stanza)
+        for stanza in found_stanzas:
+            self.iterables.remove(stanza)
+            self.xml.remove(stanza.xml)