Add callback for determining if a reply is allowed

nbxmpp/modules/entity_time.py

@@ -17,6 +17,9 @@
 
 from nbxmpp.protocol import Iq
 from nbxmpp.protocol import NodeProcessed
+from nbxmpp.protocol import Error
+from nbxmpp.protocol import ERR_FORBIDDEN
+from nbxmpp.protocol import ERR_SERVICE_UNAVAILABLE
 from nbxmpp.namespaces import Namespace
 from nbxmpp.task import iq_request_task
 from nbxmpp.structs import StanzaHandler
@@ -37,12 +40,13 @@ class EntityTime(BaseModule):
         self.handlers = [
             StanzaHandler(name='iq',
                           callback=self._answer_request,
-                          priority=30,
+                          priority=60,
                           typ='get',
                           ns=Namespace.TIME),
         ]
 
         self._enabled = False
+        self._allow_reply_func = None
 
     def disable(self):
         self._enabled = False
@@ -50,6 +54,9 @@ class EntityTime(BaseModule):
     def enable(self):
         self._enabled = True
 
+    def set_allow_reply_func(self, func):
+        self._allow_reply_func = func
+
     @iq_request_task
     def request_entity_time(self, jid):
         _task = yield
@@ -63,7 +70,13 @@ class EntityTime(BaseModule):
     def _answer_request(self, _con, stanza, _properties):
         self._log.info('Request received from %s', stanza.getFrom())
         if not self._enabled:
-            return
+            self._client.send_stanza(Error(stanza, ERR_SERVICE_UNAVAILABLE))
+            raise NodeProcessed
+
+        if self._allow_reply_func is not None:
+            if not self._allow_reply_func(stanza.getFrom()):
+                self._client.send_stanza(Error(stanza, ERR_FORBIDDEN))
+                raise NodeProcessed
 
         time, tzo = get_local_time()
         iq = stanza.buildSimpleReply('result')

nbxmpp/modules/last_activity.py

@@ -17,6 +17,9 @@
 
 from nbxmpp.protocol import Iq
 from nbxmpp.protocol import NodeProcessed
+from nbxmpp.protocol import Error
+from nbxmpp.protocol import ERR_SERVICE_UNAVAILABLE
+from nbxmpp.protocol import ERR_FORBIDDEN
 from nbxmpp.namespaces import Namespace
 from nbxmpp.task import iq_request_task
 from nbxmpp.structs import LastActivityData
@@ -34,12 +37,13 @@ class LastActivity(BaseModule):
         self.handlers = [
             StanzaHandler(name='iq',
                           callback=self._answer_request,
-                          priority=30,
+                          priority=60,
                           typ='get',
                           ns=Namespace.LAST),
         ]
 
         self._idle_func = None
+        self._allow_reply_func = None
 
     def disable(self):
         self._idle_func = None
@@ -47,6 +51,9 @@ class LastActivity(BaseModule):
     def set_idle_func(self, func):
         self._idle_func = func
 
+    def set_allow_reply_func(self, func):
+        self._allow_reply_func = func
+
     @iq_request_task
     def request_last_activity(self, jid):
         _task = yield
@@ -60,7 +67,13 @@ class LastActivity(BaseModule):
     def _answer_request(self, _client, stanza, _properties):
         self._log.info('Request received from %s', stanza.getFrom())
         if self._idle_func is None:
-            return
+            self._client.send_stanza(Error(stanza, ERR_SERVICE_UNAVAILABLE))
+            raise NodeProcessed
+
+        if self._allow_reply_func is not None:
+            if not self._allow_reply_func(stanza.getFrom()):
+                self._client.send_stanza(Error(stanza, ERR_FORBIDDEN))
+                raise NodeProcessed
 
         seconds = self._idle_func()
         iq = stanza.buildReply('result')

nbxmpp/modules/software_version.py

@@ -17,7 +17,8 @@
 
 from nbxmpp.namespaces import Namespace
 from nbxmpp.protocol import Iq
-from nbxmpp.protocol import ErrorNode
+from nbxmpp.protocol import Error
+from nbxmpp.protocol import ERR_FORBIDDEN
 from nbxmpp.protocol import NodeProcessed
 from nbxmpp.protocol import ERR_SERVICE_UNAVAILABLE
 from nbxmpp.structs import SoftwareVersionResult
@@ -37,6 +38,7 @@ class SoftwareVersion(BaseModule):
             StanzaHandler(name='iq',
                           callback=self._answer_request,
                           typ='get',
+                          priority=60,
                           ns=Namespace.VERSION),
         ]
 
@@ -45,10 +47,14 @@ class SoftwareVersion(BaseModule):
         self._os = None
 
         self._enabled = False
+        self._allow_reply_func = None
 
     def disable(self):
         self._enabled = False
 
+    def set_allow_reply_func(self, func):
+        self._allow_reply_func = func
+
     @iq_request_task
     def request_software_version(self, jid):
         _task = yield
@@ -70,19 +76,22 @@ class SoftwareVersion(BaseModule):
         if (not self._enabled or
                 self._name is None or
                 self._version is None):
-            iq = stanza.buildReply('error')
-            iq.addChild(node=ErrorNode(ERR_SERVICE_UNAVAILABLE))
-            self._log.info('Send service-unavailable')
-
-        else:
-            iq = stanza.buildReply('result')
-            query = iq.getQuery()
-            query.setTagData('name', self._name)
-            query.setTagData('version', self._version)
-            if self._os is not None:
-                query.setTagData('os', self._os)
-            self._log.info('Send software version: %s %s %s',
-                           self._name, self._version, self._os)
+            self._client.send_stanza(Error(stanza, ERR_SERVICE_UNAVAILABLE))
+            raise NodeProcessed
+
+        if self._allow_reply_func is not None:
+            if not self._allow_reply_func(stanza.getFrom()):
+                self._client.send_stanza(Error(stanza, ERR_FORBIDDEN))
+                raise NodeProcessed
+
+        iq = stanza.buildReply('result')
+        query = iq.getQuery()
+        query.setTagData('name', self._name)
+        query.setTagData('version', self._version)
+        if self._os is not None:
+            query.setTagData('os', self._os)
+        self._log.info('Send software version: %s %s %s',
+                       self._name, self._version, self._os)
 
         self._client.send_stanza(iq)
         raise NodeProcessed