Merge pull request #29 from onaio/issue-27-oauth-redirect_url

Redirect after login

README.md

@@ -83,6 +83,12 @@ OAUTH_PROVIDERS = [
 CUSTOM_SECURITY_MANAGER = CustomSecurityManager
 ```
 
+As an alternative, you can also simply add a redirect variable to the url so as to redirect after logging in. This can be added as below
+
+```
+'https://example.com/login/provider?redirect=/superset/dashboard/3/'
+```
+
 ### Add custom roles
 
 This feature allows you to add custom roles to Superset on initialization.  This is useful when you want to add custom roles to Superset during an automated deployment.

superset_patchup/oauth.py

@@ -2,7 +2,7 @@
 import logging
 import re
 
-from flask import abort, flash, redirect, request
+from flask import abort, flash, g, redirect, request, session, url_for
 
 from flask_appbuilder._compat import as_unicode
 from flask_appbuilder.security.sqla import models as ab_models
@@ -12,13 +12,73 @@
 
 from superset.security import SupersetSecurityManager
 
+import jwt
 from flask_login import login_user
 
 from superset_patchup.utils import is_safe_url
 
 
 class AuthOAuthView(SupersetAuthOAuthView):
     """ Flask-AppBuilder's Authentication OAuth view"""
+    login_template = "appbuilder/general/security/login_oauth.html"
+
+    @expose("/login/")
+    @expose("/login/<provider>")
+    @expose("/login/<provider>/<register>")
+    # pylint: disable=logging-fstring-interpolation
+    def login(self, provider=None, register=None):
+        """The login view from AuthOAuthView"""
+        logging.debug(f"Provider: {provider}")
+
+        # handle redirect
+        redirect_url = self.appbuilder.get_url_for_index
+        if request.args.get('redirect_url') is not None:
+            redirect_url = request.args.get('redirect_url')
+            if not is_safe_url(redirect_url):
+                return abort(400)
+
+        if g.user is not None and g.user.is_authenticated:
+            logging.debug(f"Already authenticated {g.user}")
+            return redirect(redirect_url)
+
+        if provider is None:
+            return self.render_template(
+                self.login_template,
+                providers=self.appbuilder.sm.oauth_providers,
+                title=self.title,
+                appbuilder=self.appbuilder,
+            )
+        logging.debug(f"Going to call authorize for: {provider}")
+        state = jwt.encode(
+            request.args.to_dict(flat=False),
+            self.appbuilder.app.config["SECRET_KEY"],
+            algorithm="HS256",
+        )
+        try:
+            if register:
+                logging.debug("Login to Register")
+                session["register"] = True
+            if provider == "twitter":
+                return self.appbuilder.sm.oauth_remotes[provider].authorize(
+                    callback=url_for(
+                        ".oauth_authorized",
+                        provider=provider,
+                        _external=True,
+                        state=state,
+                    )
+                )
+            return self.appbuilder.sm.oauth_remotes[provider].authorize(
+                callback=url_for(
+                    ".oauth_authorized",
+                    provider=provider,
+                    _external=True
+                ),
+                state=state,
+            )
+        except Exception as err:  # pylint: disable=broad-except
+            logging.error(f"Error on OAuth authorize: {err}")
+            flash(as_unicode(self.invalid_login_message), "warning")
+            return redirect(redirect_url)
 
     @expose("/oauth-authorized/<provider>")
     # pylint: disable=too-many-branches

tests/oauth/test_oauth.py

@@ -298,3 +298,30 @@ def test_oauth_authorized(
             "onadata", {"access_token": "cZpwCzYjpzuSqzekM"})
         assert mock_login.call_count == 1
         mock_redirect.assert_called_once_with("http://example.com")
+
+    @patch("superset_patchup.oauth.redirect")
+    @patch("superset_patchup.oauth.g")
+    @patch("superset_patchup.oauth.is_safe_url")
+    @patch("superset_patchup.oauth.request.args.get")
+    @patch("superset_patchup.oauth.request")
+    def test_login_redirec(
+            self,
+            mock_request,
+            mock_redirect_arg,
+            mock_safe_url,
+            mock_g,
+            mock_redirect
+    ):
+        """
+        Test that we are redirected to the redirect url when it is passed
+        as an argument to /login
+        """
+        oauth_view = AuthOAuthView()
+        oauth_view.appbuilder = MagicMock()
+
+        mock_redirect_arg.return_value = "/superset/dashboard/3"
+        mock_safe_url.return_value = True
+        mock_g.user.is_authenticated.return_value = True
+
+        oauth_view.login(provider="onadata")
+        mock_redirect.assert_called_once_with("/superset/dashboard/3")