Register: Refactor module

- Use tasks

6 files changed, 252 insertions, 210 deletions

diff --git a/nbxmpp/errors.py b/nbxmpp/errors.py
index d6131e3..c926c06 100644
--- a/nbxmpp/errors.py
+++ b/nbxmpp/errors.py
@@ -135,3 +135,21 @@ class CancelledError(BaseError):
     def __init__(self):
         BaseError.__init__(self, is_fatal=True)
         self.text = 'Task has been cancelled'
+
+
+class RegisterStanzaError(StanzaError):
+    def __init__(self, stanza, data):
+        StanzaError.__init__(self, stanza)
+        self._data = data
+
+    def get_data(self):
+        return self._data
+
+
+class ChangePasswordStanzaError(StanzaError):
+    def __init__(self, stanza, form):
+        StanzaError.__init__(self, stanza)
+        self._form = form
+
+    def get_form(self):
+        return self._form
diff --git a/nbxmpp/modules/register.py b/nbxmpp/modules/register.py
deleted file mode 100644
index 4932be7..0000000
--- a/nbxmpp/modules/register.py
+++ /dev/null
@@ -1,207 +0,0 @@
-# Copyright (C) 2018 Philipp Hörist <philipp AT hoerist.com>
-#
-# This file is part of nbxmpp.
-#
-# This program is free software; you can redistribute it and/or
-# modify it under the terms of the GNU General Public License
-# as published by the Free Software Foundation; either version 3
-# of the License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# 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 nbxmpp.namespaces import Namespace
-from nbxmpp.protocol import Iq
-from nbxmpp.protocol import isResultNode
-from nbxmpp.structs import CommonResult
-from nbxmpp.structs import CommonError
-from nbxmpp.structs import RegisterData
-from nbxmpp.structs import ChangePasswordResult
-from nbxmpp.util import call_on_response
-from nbxmpp.util import callback
-from nbxmpp.util import raise_error
-from nbxmpp.util import get_form
-from nbxmpp.const import REGISTER_FIELDS
-from nbxmpp.modules.bits_of_binary import parse_bob_data
-from nbxmpp.modules.dataforms import extend_form
-from nbxmpp.modules.dataforms import create_field
-from nbxmpp.modules.dataforms import SimpleDataForm
-from nbxmpp.modules.base import BaseModule
-
-
-class Register(BaseModule):
-    def __init__(self, client):
-        BaseModule.__init__(self, client)
-
-        self._client = client
-        self.handlers = []
-
-    @call_on_response('_default_response')
-    def unregister(self, jid=None):
-        iq = Iq('set', to=jid)
-        query = iq.setQuery()
-        query.setNamespace(Namespace.REGISTER)
-        query.addChild('remove')
-        return iq
-
-    @call_on_response('_on_register_form')
-    def request_register_form(self, jid=None):
-        if jid is None:
-            jid = self._client.domain
-        return Iq('get', Namespace.REGISTER, to=jid)
-
-    @callback
-    def _on_register_form(self, stanza):
-        if not isResultNode(stanza):
-            return raise_error(self._log.info, stanza)
-
-        query = stanza.getQuery()
-        instructions = query.getTagData('instructions') or None
-
-        data = RegisterData(instructions=instructions,
-                            form=self._parse_form(stanza),
-                            fields_form=self._parse_fields_form(query),
-                            oob_url=self._parse_oob_url(query),
-                            bob_data=parse_bob_data(query))
-
-        if (data.form is None and
-                data.fields_form is None and
-                data.oob_url is None):
-            return raise_error(self._log.info, stanza, 'stanza-malformed')
-        return data
-
-    @call_on_response('_on_submit_result')
-    def submit_register_form(self, jid, form):
-        if jid is None:
-            jid = self._client.domain
-
-        iq = Iq('set', Namespace.REGISTER, to=jid)
-
-        if form.is_fake_form():
-            query = iq.getTag('query')
-            for field in form.iter_fields():
-                if field.var == 'fakeform':
-                    continue
-                query.addChild(field.var, payload=[field.value])
-            return iq
-
-        iq.setQueryPayload(form)
-        return iq
-
-    @callback
-    def _on_submit_result(self, stanza):
-        if isResultNode(stanza):
-            return CommonResult(jid=stanza.getFrom())
-
-        query = stanza.getTag('query', namespace=Namespace.REGISTER)
-        if query is None:
-            return RegisterError(stanza, None)
-
-        instructions = query.getTagData('instructions') or None
-
-        data = RegisterData(instructions=instructions,
-                            form=self._parse_form(stanza),
-                            fields_form=self._parse_fields_form(query),
-                            oob_url=self._parse_oob_url(query),
-                            bob_data=parse_bob_data(query))
-
-        return RegisterError(stanza, data)
-
-    @callback
-    def _default_response(self, stanza):
-        if not isResultNode(stanza):
-            return raise_error(self._log.info, stanza)
-        return CommonResult(jid=stanza.getFrom())
-
-    @staticmethod
-    def _parse_oob_url(query):
-        oob = query.getTag('x', namespace=Namespace.X_OOB)
-        if oob is not None:
-            return oob.getTagData('url') or None
-        return None
-
-    def _parse_form(self, stanza):
-        query = stanza.getTag('query', namespace=Namespace.REGISTER)
-        form = query.getTag('x', namespace=Namespace.DATA)
-        if form is None:
-            return None
-
-        form = extend_form(node=form)
-        field = form.vars.get('FORM_TYPE')
-        if field is None:
-            self._log.warning('No FORM_TYPE found')
-            self._log.warning(stanza)
-            return None
-
-        # Invalid urn:xmpp:captcha used by ejabberd
-        # See https://github.com/processone/ejabberd/issues/3093
-        if field.value in ('jabber:iq:register', 'urn:xmpp:captcha'):
-            return form
-        return None
-
-    @staticmethod
-    def _parse_fields_form(query):
-        fields = []
-        for field in query.getChildren():
-            field_name = field.getName()
-            if field_name not in REGISTER_FIELDS:
-                continue
-
-            required = field_name in ('username', 'password')
-            typ = 'text-single' if field_name != 'password' else 'text-private'
-            fields.append(create_field(typ=typ,
-                                       var=field_name,
-                                       required=required))
-
-        if not fields:
-            return None
-
-        fields.append(create_field(typ='hidden', var='fakeform'))
-        return SimpleDataForm(type_='form',
-                              instructions=query.getTagData('instructions'),
-                              fields=fields)
-
-    @call_on_response('_on_password_change')
-    def change_password(self, password):
-        domain = self._client.get_bound_jid().domain
-        username = self._client.get_bound_jid().localpart
-        iq = Iq('set', Namespace.REGISTER, to=domain)
-        query = iq.getQuery()
-        query.setTagData('username', username)
-        query.setTagData('password', password)
-        return iq
-
-    @callback
-    def _on_password_change(self, stanza):
-        if isResultNode(stanza):
-            return ChangePasswordResult(successful=True)
-
-        if stanza.getQuery() is None:
-            return raise_error(self._log.info, stanza)
-
-        form = get_form(stanza.getQuery(),
-                        'jabber:iq:register:changepassword')
-        if form is None:
-            return raise_error(self._log.info, stanza)
-        return ChangePasswordResult(successful=False, form=form)
-
-    @call_on_response('_default_response')
-    def change_password_with_form(self, form):
-        domain = self._client.get_bound_jid().domain
-        iq = Iq('set', Namespace.REGISTER, to=domain)
-        iq.setQueryPayload(form)
-        return iq
-
-
-class RegisterError(CommonError):
-    def __init__(self, stanza, data):
-        CommonError.__init__(self, stanza)
-        self._data = data
-
-    def get_data(self):
-        return self._data
diff --git a/nbxmpp/modules/register/__init__.py b/nbxmpp/modules/register/__init__.py
new file mode 100644
index 0000000..da1db62
--- /dev/null
+++ b/nbxmpp/modules/register/__init__.py
@@ -0,0 +1 @@
+from .register import Register
diff --git a/nbxmpp/modules/register/register.py b/nbxmpp/modules/register/register.py
new file mode 100644
index 0000000..2ba402c
--- /dev/null
+++ b/nbxmpp/modules/register/register.py
@@ -0,0 +1,102 @@
+# Copyright (C) 2020 Philipp Hörist <philipp AT hoerist.com>
+#
+# This file is part of nbxmpp.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 3
+# of the License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# 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 nbxmpp.namespaces import Namespace
+from nbxmpp.protocol import Iq
+from nbxmpp.util import get_form
+from nbxmpp.task import iq_request_task
+from nbxmpp.errors import StanzaError
+from nbxmpp.errors import RegisterStanzaError
+from nbxmpp.errors import ChangePasswordStanzaError
+from nbxmpp.modules.base import BaseModule
+from nbxmpp.modules.util import process_response
+from nbxmpp.modules.register.util import _make_unregister_request
+from nbxmpp.modules.register.util import _make_register_form
+from nbxmpp.modules.register.util import _make_password_change_request
+from nbxmpp.modules.register.util import _make_password_change_with_form
+from nbxmpp.modules.register.util import _parse_register_data
+
+
+class Register(BaseModule):
+    def __init__(self, client):
+        BaseModule.__init__(self, client)
+
+        self._client = client
+        self.handlers = []
+
+    @iq_request_task
+    def unregister(self, jid=None):
+        _task = yield
+
+        response = yield _make_unregister_request(jid)
+        yield process_response(response)
+
+    @iq_request_task
+    def request_register_form(self, jid=None):
+        _task = yield
+
+        if jid is None:
+            jid = self._client.domain
+
+        response = yield Iq('get', Namespace.REGISTER, to=jid)
+        if response.isError():
+            raise StanzaError(response)
+
+        yield _parse_register_data(response)
+
+    @iq_request_task
+    def submit_register_form(self, form, jid=None):
+        _task = yield
+
+        if jid is None:
+            jid = self._client.domain
+
+        response = yield _make_register_form(jid, form)
+        if not response.isError():
+            yield process_response(response)
+
+        else:
+            data = _parse_register_data(response)
+            raise RegisterStanzaError(response, data)
+
+    @iq_request_task
+    def change_password(self, password):
+        _task = yield
+
+        response = yield _make_password_change_request(
+            self._client.domain, self._client.username, password)
+        if not response.isError():
+            yield process_response(response)
+
+        else:
+            query = response.getQuery()
+            if query is None:
+                raise StanzaError(response)
+
+            form = get_form(query, 'jabber:iq:register:changepassword')
+            if form is None or response.getType() != 'modify':
+                raise StanzaError(response)
+
+            raise ChangePasswordStanzaError(response, form)
+
+    @iq_request_task
+    def change_password_with_form(self, form):
+        _task = yield
+
+        response = yield _make_password_change_with_form(self._client.domain,
+                                                         form)
+        yield process_response(response)
diff --git a/nbxmpp/modules/register/util.py b/nbxmpp/modules/register/util.py
new file mode 100644
index 0000000..53f5dd0
--- /dev/null
+++ b/nbxmpp/modules/register/util.py
@@ -0,0 +1,131 @@
+# Copyright (C) 2018 Philipp Hörist <philipp AT hoerist.com>
+#
+# This file is part of nbxmpp.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 3
+# of the License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# 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 nbxmpp.namespaces import Namespace
+from nbxmpp.protocol import Iq
+from nbxmpp.const import REGISTER_FIELDS
+from nbxmpp.structs import RegisterData
+from nbxmpp.errors import StanzaError
+from nbxmpp.errors import MalformedStanzaError
+from nbxmpp.modules.dataforms import create_field
+from nbxmpp.modules.dataforms import extend_form
+from nbxmpp.modules.dataforms import SimpleDataForm
+from nbxmpp.modules.bits_of_binary import parse_bob_data
+
+
+def _make_password_change_request(domain, username, password):
+    iq = Iq('set', Namespace.REGISTER, to=domain)
+    query = iq.getQuery()
+    query.setTagData('username', username)
+    query.setTagData('password', password)
+    return iq
+
+
+def _make_password_change_with_form(domain, form):
+    iq = Iq('set', Namespace.REGISTER, to=domain)
+    iq.setQueryPayload(form)
+    return iq
+
+
+def _make_register_form(jid, form):
+    iq = Iq('set', Namespace.REGISTER, to=jid)
+    if form.is_fake_form():
+        query = iq.getTag('query')
+        for field in form.iter_fields():
+            if field.var == 'fakeform':
+                continue
+            query.addChild(field.var, payload=[field.value])
+        return iq
+
+    iq.setQueryPayload(form)
+    return iq
+
+
+def _make_unregister_request(jid):
+    iq = Iq('set', to=jid)
+    query = iq.setQuery()
+    query.setNamespace(Namespace.REGISTER)
+    query.addChild('remove')
+    return iq
+
+
+def _parse_oob_url(query):
+    oob = query.getTag('x', namespace=Namespace.X_OOB)
+    if oob is not None:
+        return oob.getTagData('url') or None
+    return None
+
+
+def _parse_form(stanza):
+    query = stanza.getTag('query', namespace=Namespace.REGISTER)
+    form = query.getTag('x', namespace=Namespace.DATA)
+    if form is None:
+        return None
+
+    form = extend_form(node=form)
+    field = form.vars.get('FORM_TYPE')
+    if field is None:
+        return None
+
+    # Invalid urn:xmpp:captcha used by ejabberd
+    # See https://github.com/processone/ejabberd/issues/3045
+    if field.value in ('jabber:iq:register', 'urn:xmpp:captcha'):
+        return form
+    return None
+
+
+def _parse_fields_form(query):
+    fields = []
+    for field in query.getChildren():
+        field_name = field.getName()
+        if field_name not in REGISTER_FIELDS:
+            continue
+
+        required = field_name in ('username', 'password')
+        typ = 'text-single' if field_name != 'password' else 'text-private'
+        fields.append(create_field(typ=typ,
+                                   var=field_name,
+                                   required=required))
+
+    if not fields:
+        return None
+
+    fields.append(create_field(typ='hidden', var='fakeform'))
+    return SimpleDataForm(type_='form',
+                          instructions=query.getTagData('instructions'),
+                          fields=fields)
+
+
+def _parse_register_data(response):
+    query = response.getTag('query', namespace=Namespace.REGISTER)
+    if query is None:
+        raise StanzaError(response)
+
+    instructions = query.getTagData('instructions') or None
+
+    data = RegisterData(instructions=instructions,
+                        form=_parse_form(response),
+                        fields_form=_parse_fields_form(query),
+                        oob_url=_parse_oob_url(query),
+                        bob_data=parse_bob_data(query))
+
+    if (data.form is None and
+            data.fields_form is None and
+            data.oob_url is None):
+        raise MalformedStanzaError('invalid register response', response)
+
+    return data
diff --git a/nbxmpp/structs.py b/nbxmpp/structs.py
index 2631336..c75c598 100644
--- a/nbxmpp/structs.py
+++ b/nbxmpp/structs.py
@@ -120,9 +120,6 @@ SecurityLabel = namedtuple('SecurityLabel', 'displaymarking')
 
 RegisterData = namedtuple('RegisterData', 'instructions form fields_form oob_url bob_data')
 
-ChangePasswordResult = namedtuple('ChangePasswordResult', 'successful form')
-ChangePasswordResult.__new__.__defaults__ = (None,)
-
 HTTPUploadData = namedtuple('HTTPUploadData', 'put_uri get_uri headers')
 HTTPUploadData.__new__.__defaults__ = (None,)