diff --git a/app/auth/views/facebook.py b/app/auth/views/facebook.py index ef5cc305..e8712572 100644 --- a/app/auth/views/facebook.py +++ b/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) diff --git a/app/auth/views/github.py b/app/auth/views/github.py index 6e336daa..93120a2f 100644 --- a/app/auth/views/github.py +++ b/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) diff --git a/app/auth/views/google.py b/app/auth/views/google.py index 775f3052..60480b1d 100644 --- a/app/auth/views/google.py +++ b/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: diff --git a/app/auth/views/login.py b/app/auth/views/login.py index 73d8a42b..4c828d6e 100644 --- a/app/auth/views/login.py +++ b/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", diff --git a/app/auth/views/login_utils.py b/app/auth/views/login_utils.py new file mode 100644 index 00000000..246e22ca --- /dev/null +++ b/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"))