Redirect user to MFA page if they enable MFA

app/auth/views/facebook.py

@@ -10,6 +10,7 @@ from app.config import URL, FACEBOOK_CLIENT_ID, FACEBOOK_CLIENT_SECRET
 from app.extensions import db
 from app.log import LOG
 from app.models import User
+from .login_utils import after_login
 
 _authorization_base_url = "https://www.facebook.com/dialog/oauth"
 _token_url = "https://graph.facebook.com/oauth/access_token"
@@ -99,7 +100,6 @@ def facebook_callback():
             user.profile_picture_id = file.id
             db.session.commit()
 
-        login_user(user)
     # create user
     else:
         LOG.d("create facebook user with %s", facebook_user_data)
@@ -116,6 +116,7 @@ def facebook_callback():
 
         flash(f"Welcome to SimpleLogin {user.name}!", "success")
 
+    next_url = None
     # The activation link contains the original page, for ex authorize page
     if "facebook_next_url" in session:
         next_url = session["facebook_next_url"]
@@ -124,7 +125,4 @@ def facebook_callback():
         # reset the next_url to avoid user getting redirected at each login :)
         session.pop("facebook_next_url", None)
 
-        return redirect(next_url)
-    else:
-        LOG.debug("redirect user to dashboard")
-        return redirect(url_for("dashboard.index"))
+    return after_login(user, next_url)

app/auth/views/github.py

@@ -1,9 +1,10 @@
-from flask import request, session, redirect, url_for, flash
+from flask import request, session, redirect, flash
 from flask_login import login_user
 from requests_oauthlib import OAuth2Session
 
 from app import email_utils
 from app.auth.base import auth_bp
+from app.auth.views.login_utils import after_login
 from app.config import GITHUB_CLIENT_ID, GITHUB_CLIENT_SECRET, URL
 from app.extensions import db
 from app.log import LOG
@@ -81,10 +82,8 @@ def github_callback():
 
     user = User.get_by(email=email)
 
-    if user:
-        login_user(user)
     # create user
-    else:
+    if not user:
         LOG.d("create github user")
         user = User.create(
             email=email, name=github_user_data.get("name") or "", activated=True
@@ -96,10 +95,6 @@ def github_callback():
         flash(f"Welcome to SimpleLogin {user.name}!", "success")
 
     # The activation link contains the original page, for ex authorize page
-    if "next" in request.args:
-        next_url = request.args.get("next")
-        LOG.debug("redirect user to %s", next_url)
-        return redirect(next_url)
-    else:
-        LOG.debug("redirect user to dashboard")
-        return redirect(url_for("dashboard.index"))
+    next_url = request.args.get("next") if request.args else None
+
+    return after_login(user, next_url)

app/auth/views/google.py

@@ -1,4 +1,4 @@
-from flask import request, session, redirect, url_for, flash
+from flask import request, session, redirect, flash
 from flask_login import login_user
 from requests_oauthlib import OAuth2Session
 
@@ -9,6 +9,7 @@ from app.extensions import db
 from app.log import LOG
 from app.models import User, File
 from app.utils import random_string
+from .login_utils import after_login
 
 _authorization_base_url = "https://accounts.google.com/o/oauth2/v2/auth"
 _token_url = "https://www.googleapis.com/oauth2/v4/token"
@@ -89,8 +90,6 @@ def google_callback():
             file = create_file_from_url(picture_url)
             user.profile_picture_id = file.id
             db.session.commit()
-
-        login_user(user)
     # create user
     else:
         LOG.d("create google user with %s", google_user_data)
@@ -107,6 +106,7 @@ def google_callback():
 
         flash(f"Welcome to SimpleLogin {user.name}!", "success")
 
+    next_url = None
     # The activation link contains the original page, for ex authorize page
     if "google_next_url" in session:
         next_url = session["google_next_url"]
@@ -115,10 +115,7 @@ def google_callback():
         # reset the next_url to avoid user getting redirected at each login :)
         session.pop("google_next_url", None)
 
-        return redirect(next_url)
-    else:
-        LOG.debug("redirect user to dashboard")
-        return redirect(url_for("dashboard.index"))
+    return after_login(user, next_url)
 
 
 def create_file_from_url(url) -> File:

app/auth/views/login.py

@@ -1,9 +1,10 @@
 from flask import request, render_template, redirect, url_for, flash
-from flask_login import login_user, current_user
+from flask_login import current_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
+from app.auth.views.login_utils import after_login
 from app.log import LOG
 from app.models import User
 
@@ -37,16 +38,7 @@ def login():
                 "error",
             )
         else:
-            LOG.debug("log user %s in", user)
-            login_user(user)
-
-            # User comes to login page from another page
-            if next_url:
-                LOG.debug("redirect user to %s", next_url)
-                return redirect(next_url)
-            else:
-                LOG.debug("redirect user to dashboard")
-                return redirect(url_for("dashboard.index"))
+            return after_login(user, next_url)
 
     return render_template(
         "auth/login.html",

app/auth/views/login_utils.py

@@ -0,0 +1,30 @@
+from flask import session, redirect, url_for
+from flask_login import login_user
+
+from app.config import MFA_USER_ID
+from app.log import LOG
+
+
+def after_login(user, next_url):
+    """
+    Redirect to the correct page after login.
+    If user enables MFA: redirect user to MFA page
+    Otherwise redirect to dashboard page if no next_url
+    """
+    if user.enable_otp:
+        session[MFA_USER_ID] = user.id
+        if next_url:
+            return redirect(url_for("auth.mfa", next_url=next_url))
+        else:
+            return redirect(url_for("auth.mfa"))
+    else:
+        LOG.debug("log user %s in", user)
+        login_user(user)
+
+        # User comes to login page from another page
+        if next_url:
+            LOG.debug("redirect user to %s", next_url)
+            return redirect(next_url)
+        else:
+            LOG.debug("redirect user to dashboard")
+            return redirect(url_for("dashboard.index"))