Add registration to components

1 files changed, 152 insertions, 4 deletions

diff --git a/slixmpp/plugins/xep_0077/register.py b/slixmpp/plugins/xep_0077/register.py
index a19f517f..953fee70 100644
--- a/slixmpp/plugins/xep_0077/register.py
+++ b/slixmpp/plugins/xep_0077/register.py
@@ -8,6 +8,8 @@ import ssl
 
 from slixmpp.stanza import StreamFeatures, Iq
 from slixmpp.xmlstream import register_stanza_plugin, JID
+from slixmpp.xmlstream.handler import CoroutineCallback
+from slixmpp.xmlstream.matcher import StanzaPath
 from slixmpp.plugins import BasePlugin
 from slixmpp.plugins.xep_0077 import stanza, Register, RegisterFeature
 
@@ -19,6 +21,36 @@ class XEP_0077(BasePlugin):
 
     """
     XEP-0077: In-Band Registration
+
+    Events:
+
+    ::
+
+        user_register           -- After succesful validation and add to the user store
+                                   in api["user_validate"]
+        user_unregister         -- After succesful user removal in api["user_remove"]
+    
+    Config:
+
+    ::
+
+        form_fields are only form_instructions are only used for component registration
+        in case api["make_registration_form"] is not overriden.
+
+    API:
+    
+    ::
+
+        user_get(jid, node, ifrom, iq)
+            Returns a dict-like object containing `form_fields` for this user or None
+        user_remove(jid, node, ifrom, iq)
+            Removes a user or raise KeyError in case the user is not found in the user store
+        make_registration_form(self, jid, node, ifrom, iq)
+            Returns an iq reply to a registration form request, pre-filled and with
+            <registered/> in case the requesting entity is already registered to us
+        user_validate((self, jid, node, ifrom, registration)
+            Add the user to the user store or raise ValueError(msg) if any problem is encountered
+            msg is sent back to the XMPP client as an error message.
     """
 
     name = 'xep_0077'
@@ -28,18 +60,36 @@ class XEP_0077(BasePlugin):
     default_config = {
         'create_account': True,
         'force_registration': False,
-        'order': 50
+        'order': 50,
+        "form_fields": {"username", "password"},
+        "form_instructions": "Enter your credentials",
     }
 
     def plugin_init(self):
         register_stanza_plugin(StreamFeatures, RegisterFeature)
         register_stanza_plugin(Iq, Register)
 
-        if not self.xmpp.is_component:
-            self.xmpp.register_feature('register',
+        if self.xmpp.is_component:
+            self.xmpp["xep_0030"].add_feature("jabber:iq:register")
+            self.xmpp.register_handler(
+                CoroutineCallback(
+                    "registration",
+                    StanzaPath("/iq/register"),
+                    self._handle_registration,
+                )
+            )
+            self._user_store = {}
+            self.api.register(self._user_get, "user_get")
+            self.api.register(self._user_remove, "user_remove")
+            self.api.register(self._make_registration_form, "make_registration_form")
+            self.api.register(self._user_validate, "user_validate")
+        else:
+            self.xmpp.register_feature(
+                "register",
                 self._handle_register_feature,
                 restart=False,
-                order=self.order)
+                order=self.order,
+            )
 
         register_stanza_plugin(Register, self.xmpp['xep_0004'].stanza.Form)
         register_stanza_plugin(Register, self.xmpp['xep_0066'].stanza.OOB)
@@ -50,6 +100,92 @@ class XEP_0077(BasePlugin):
         if not self.xmpp.is_component:
             self.xmpp.unregister_feature('register', self.order)
 
+    def _user_get(self, jid, node, ifrom, iq):
+        return self._user_store.get(iq["from"].bare)
+    
+    def _user_remove(self, jid, node, ifrom, iq):
+        return self._user_store.pop(iq["from"].bare)
+
+    def _make_registration_form(self, jid, node, ifrom, iq: Iq):
+        reg = iq["register"]
+        user = self.api["user_get"](None, None, None, iq)
+
+
+        if user is None:
+            user = {}
+        else:
+            reg["registered"] = True
+
+        reg["instructions"] = self.form_instructions
+
+        for field in self.form_fields:
+            data = user.get(field, "")
+            if data:
+                reg[field] = data
+            else:
+                # Add a blank field
+                reg.add_field(field)
+
+        reply = iq.reply()
+        reply.set_payload(reg.xml)
+        return reply
+
+    def _user_validate(self, jid, node, ifrom, registration):
+        self._user_store[ifrom.bare] = {key: registration[key] for key in self.form_fields}
+
+    async def _handle_registration(self, iq: Iq):
+        if iq["type"] == "get":
+            self._send_form(iq)
+        elif iq["type"] == "set":
+            if iq["register"]["remove"]:
+                try:
+                    self.api["user_remove"](None, None, iq["from"], iq)
+                except KeyError:
+                    _send_error(
+                        iq,
+                        "404",
+                        "cancel",
+                        "item-not-found",
+                        "User not found",
+                    )
+                else:
+                    reply = iq.reply()
+                    reply.send()
+                    self.xmpp.event("user_unregister", iq)
+                return
+
+
+            for field in self.form_fields:
+                if not iq["register"][field]:
+                    # Incomplete Registration
+                    _send_error(
+                        iq,
+                        "406",
+                        "modify",
+                        "not-acceptable",
+                        "Please fill in all fields.",
+                    )
+                    return
+
+            try:
+                self.api["user_validate"](None, None, iq["from"], iq["register"])
+            except ValueError as e:
+                _send_error(
+                    iq,
+                    "406",
+                    "modify",
+                    "not-acceptable",
+                    e.args,
+                )
+            else:
+                reply = iq.reply()
+                reply.send()
+                self.xmpp.event("user_register", iq)
+
+    def _send_form(self, iq):
+        reply = self.api["make_registration_form"](None, None, iq["from"], iq)
+        reply.send()
+
     def _force_registration(self, event):
         if self.force_registration:
             self.xmpp.add_filter('in', self._force_stream_feature)
@@ -109,3 +245,15 @@ class XEP_0077(BasePlugin):
             iq['register']['username'] = self.xmpp.boundjid.user
         iq['register']['password'] = password
         return iq.send(timeout=timeout, callback=callback)
+
+def _send_error(iq, code, error_type, name, text=""):
+    # It would be nice to raise XMPPError but the iq payload
+    # should include the register info
+    reply = iq.reply()
+    reply.set_payload(iq["register"].xml)
+    reply.error()
+    reply["error"]["code"] = code
+    reply["error"]["type"] = error_type
+    reply["error"]["condition"] = name
+    reply["error"]["text"] = text
+    reply.send()