Roster: Refactor roster request

nbxmpp/modules/roster.py

@@ -25,6 +25,7 @@ from nbxmpp.structs import RosterData
 from nbxmpp.structs import RosterItem
 # from nbxmpp.structs import StanzaHandler
 from nbxmpp.errors import StanzaError
+from nbxmpp.errors import MalformedStanzaError
 from nbxmpp.task import iq_request_task
 from nbxmpp.modules.base import BaseModule
 from nbxmpp.modules.util import process_response
@@ -47,12 +48,27 @@ class Roster(BaseModule):
     def request_roster(self, version=None):
         _task = yield
 
-        self._log.info('Request Roster, version: %s', version)
-        response = yield _make_request(version)
+        ver_support = self._client.features.has_roster_version()
+        if not ver_support:
+            version = None
+
+        if ver_support and version is None:
+            version = ''
+
+        self._log.info('Request Roster, version: %s, ver support: %s',
+                       version, ver_support)
+
+        response = yield _make_request(version, ver_support)
         if response.isError():
             raise StanzaError(response)
 
-        yield _parse_push(response)
+        query = response.getTag('query', namespace=Namespace.ROSTER)
+        if query is None:
+            if not ver_support:
+                raise MalformedStanzaError('query node missing', response)
+            yield RosterData(None, version)
+
+        yield self._parse_push(response, ver_support)
 
     def _process_roster_push(self, _client, stanza, properties):
         from_ = stanza.getFrom()
@@ -61,7 +77,8 @@ class Roster(BaseModule):
                 self._log.warning('Malicious Roster Push from %s', from_)
                 raise NodeProcessed
 
-        properties.roster = _parse_push(stanza)
+        ver_support = self._client.features.has_roster_version()
+        properties.roster = self._parse_push(stanza, ver_support)
 
         self._log.info('Roster Push, version: %s', properties.roster.version)
         for item in properties.roster.items:
@@ -90,6 +107,29 @@ class Roster(BaseModule):
         response = yield _make_set(jid, name, groups)
         yield process_response(response)
 
+    def _parse_push(self, stanza, ver_support):
+        query = stanza.getTag('query', namespace=Namespace.ROSTER)
+
+        version = None
+        if ver_support:
+            version = query.getAttr('ver')
+            if version is None:
+                raise MalformedStanzaError('ver attribute missing', stanza)
+
+        pushed_items = []
+        for item in query.getTags('item'):
+            try:
+                roster_item = RosterItem.from_node(item)
+            except Exception:
+                self._log.warning('Invalid roster item')
+                self._log.warning(stanza)
+                continue
+
+            pushed_items.append(roster_item)
+
+        return RosterData(pushed_items, version)
+
+
 
 def _make_delete(jid):
     return Iq('set',
@@ -112,36 +152,11 @@ def _make_set(jid, name, groups=None):
     return iq
 
 
-def _make_request(version):
+def _make_request(version, roster_ver_support):
     iq = Iq('get', Namespace.ROSTER)
-    if version is not None:
+    if version is None:
+        version = ''
+
+    if roster_ver_support:
         iq.setTagAttr('query', 'ver', version)
     return iq
-
-
-def _get_item_attrs(item):
-    default_attrs = {'name': None,
-                     'ask': None,
-                     'subscription': None,
-                     'groups': []}
-
-    attrs = item.getAttrs()
-    del attrs['jid']
-    groups = {group.getData() for group in item.getTags('group')}
-    attrs['groups'] = list(groups)
-
-    default_attrs.update(attrs)
-    return default_attrs
-
-
-def _parse_push(stanza):
-    query = stanza.getTag('query', namespace=Namespace.ROSTER)
-    version = query.getAttr('ver')
-    pushed_items = []
-
-    for item in query.getTags('item'):
-        jid = JID.from_string(item.getAttr('jid'))
-        attrs = _get_item_attrs(item)
-        pushed_items.append(RosterItem(attrs, jid))
-
-    return RosterData(pushed_items, version)

nbxmpp/structs.py

@@ -15,9 +15,13 @@
 # You should have received a copy of the GNU General Public License
 # along with this program; If not, see <http://www.gnu.org/licenses/>.
 
+from typing import List
+
 import time
 import random
 from collections import namedtuple
+from dataclasses import dataclass
+from dataclasses import field
 
 from gi.repository import Soup
 from gi.repository import Gio
@@ -25,6 +29,7 @@ from gi.repository import Gio
 from nbxmpp.namespaces import Namespace
 from nbxmpp.protocol import Protocol
 from nbxmpp.protocol import Node
+from nbxmpp.protocol import JID
 from nbxmpp.const import MessageType
 from nbxmpp.const import AvatarState
 from nbxmpp.const import StatusCode
@@ -126,10 +131,40 @@ MAMQueryData = namedtuple('MAMQueryData', 'jid rsm complete')
 
 MAMPreferencesData = namedtuple('MAMPreferencesData', 'default always never')
 
+LastActivityData = namedtuple('LastActivityData', 'seconds status')
+
 RosterData = namedtuple('RosterData', 'items version')
-RosterItem = namedtuple('RosterItem', 'data jid')
 
-LastActivityData = namedtuple('LastActivityData', 'seconds status')
+
+@dataclass
+class RosterItem:
+    jid: JID
+    name: str = None
+    ask: str = None
+    subscription: str = None
+    groups: List[str] = field(default_factory=list)
+
+    @classmethod
+    def from_node(cls, node):
+        attrs = node.getAttrs()
+        jid = attrs.get('jid')
+        if jid is None:
+            raise Exception('jid attribute missing')
+
+        jid = JID.from_string(jid)
+        attrs['jid'] = jid
+
+        groups = {group.getData() for group in node.getTags('group')}
+        attrs['groups'] = list(groups)
+
+        return cls(**attrs)
+
+    def asdict(self):
+        return {'jid': self.jid,
+                'name': self.name,
+                'ask': self.ask,
+                'subscription': self.subscription,
+                'groups': self.groups}
 
 
 class DiscoInfo(namedtuple('DiscoInfo', 'stanza identities features dataforms timestamp')):