Feature: redirect after login/register with OAuth

flask_appbuilder/security/decorators.py

@@ -1,7 +1,7 @@
 import logging
 import functools
 
-from flask import flash, redirect, url_for, make_response, jsonify
+from flask import flash, redirect, url_for, make_response, jsonify, request
 from .._compat import as_unicode
 from ..const import LOGMSG_ERR_SEC_ACCESS_DENIED, FLAMSG_ERR_SEC_ACCESS_DENIED, PERMISSION_PREFIX
 
@@ -27,7 +27,7 @@ def wraps(self, *args, **kwargs):
         else:
             log.warning(LOGMSG_ERR_SEC_ACCESS_DENIED.format(permission_str, self.__class__.__name__))
             flash(as_unicode(FLAMSG_ERR_SEC_ACCESS_DENIED), "danger")
-        return redirect(url_for(self.appbuilder.sm.auth_view.__class__.__name__ + ".login"))
+        return redirect(url_for(self.appbuilder.sm.auth_view.__class__.__name__ + ".login", next=request.url))
     f._permission_name = permission_str
     return functools.update_wrapper(wraps, f)
 
@@ -56,7 +56,7 @@ def wraps(self, *args, **kwargs):
                                               'severity': 'danger'}), 401)
             response.headers['Content-Type'] = "application/json"
             return response
-        return redirect(url_for(self.appbuilder.sm.auth_view.__class__.__name__ + ".login"))
+        return redirect(url_for(self.appbuilder.sm.auth_view.__class__.__name__ + ".login", next=request.url))
     f._permission_name = permission_str
     return functools.update_wrapper(wraps, f)
 

flask_appbuilder/security/manager.py

@@ -839,9 +839,9 @@ def auth_user_oauth(self, userinfo):
         if not user:
             user = self.add_user(
                     username=userinfo['username'],
-                    first_name=userinfo['first_name'],
-                    last_name=userinfo['last_name'],
-                    email=userinfo['email'],
+                    first_name=userinfo.get('first_name', ''),
+                    last_name=userinfo.get('last_name', ''),
+                    email=userinfo.get('email', ''),
                     role=self.find_role(self.auth_user_registration_role)
                 )
             if not user:

flask_appbuilder/security/views.py

@@ -7,6 +7,7 @@
 from wtforms.validators import EqualTo
 from flask_babel import lazy_gettext
 from flask_login import login_user, logout_user
+import jwt
 
 from ..views import ModelView, SimpleFormView, expose
 from ..baseviews import BaseView
@@ -503,11 +504,23 @@ def login(self, provider=None, register=None):
                                appbuilder=self.appbuilder)
         else:
             log.debug("Going to call authorize for: {0}".format(provider))
+            state = jwt.encode(
+                request.args.to_dict(flat=False),
+                self.appbuilder.app.config['SECRET_KEY'],
+                algorithm='HS256')
             try:
                 if register:
                     log.debug('Login to Register')
                     session['register'] = True
-                return self.appbuilder.sm.oauth_remotes[provider].authorize(callback=url_for('.oauth_authorized',provider=provider, _external=True))
+                if provider == 'twitter':
+                    return self.appbuilder.sm.oauth_remotes[provider].authorize(
+                        callback=url_for
+                            ('.oauth_authorized', provider=provider, _external=True, state=state))
+                else:
+                    return self.appbuilder.sm.oauth_remotes[provider].authorize(
+                        callback=url_for
+                            ('.oauth_authorized', provider=provider, _external=True),
+                            state=state)
             except Exception as e:
                 log.error("Error on OAuth authorize: {0}".format(e))
                 flash(as_unicode(self.invalid_login_message), 'warning')
@@ -550,7 +563,20 @@ def oauth_authorized(self, provider):
             return redirect('login')
         else:
             login_user(user)
-            return redirect(self.appbuilder.get_url_for_index)
+            try:
+                state = jwt.decode(
+                    request.args['state'],
+                    self.appbuilder.app.config['SECRET_KEY'],
+                    algorithms=['HS256'])
+            except jwt.InvalidTokenError:
+                raise AuthenticationError('State signature is not valid!')
+
+            try:
+                next_url = state['next'][0]
+            except (KeyError, IndexError):
+                next_url = self.appbuilder.get_url_for_index
+
+            return redirect(next_url)
 
 
 class AuthRemoteUserView(AuthView):

flask_appbuilder/templates/appbuilder/general/security/login_oauth.html

@@ -5,8 +5,9 @@
 
 <script type="text/javascript">
 
-var baseLoginUrl = {{url_for('AuthOAuthView.login')}};
-var baseRegisterUrl = {{url_for('AuthOAuthView.login')}};
+var baseLoginUrl = "{{url_for('AuthOAuthView.login')}}";
+var baseRegisterUrl = "{{url_for('AuthOAuthView.login')}}";
+var next = "?next={{request.args.get('next', '')}}"
 
 var currentSelection = "";
 
@@ -20,13 +21,13 @@
 
 function signin() {
     if (currentSelection != "") {
-        window.location.href = baseLoginUrl + currentSelection;
+        window.location.href = baseLoginUrl + currentSelection + next;
     }
 }
 
 function register() {
     if (currentSelection != "") {
-        window.location.href = baseRegisterUrl + currentSelection + '/register';
+        window.location.href = baseRegisterUrl + currentSelection + '/register' + next;
     }
 }
 

rtd_requirements.txt

@@ -8,3 +8,4 @@ Flask-Login>=0.3,<0.5
 Flask-OpenID>=1.2.5,<2
 Flask-SQLAlchemy>=2.3,<3
 Flask-WTF>=0.14.2,<1
+PyJWT>=1.7.1

setup.py

@@ -44,6 +44,7 @@ def desc():
         'Flask-SQLAlchemy>=2.3,<3',
         'Flask-WTF>=0.14.2,<1',
         'python-dateutil>=2.3,<3',
+        'PyJWT>=1.7.1',
     ],
     tests_require=[
         'nose>=1.0',