app-MAIL-temp/app/auth/views/login_utils.py

49 lines
1.4 KiB
Python
Raw Normal View History

2020-04-09 22:22:26 +02:00
from typing import Optional
from flask import session, redirect, url_for, request
from flask_login import login_user
from app.config import MFA_USER_ID
from app.log import LOG
2020-04-09 22:22:26 +02:00
from app.models import Referral
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"))
2020-04-09 22:22:26 +02:00
# name of the cookie that stores the referral code
_REFERRAL_COOKIE = "slref"
def get_referral() -> Optional[Referral]:
"""Get the eventual referral stored in cookie"""
# whether user arrives via a referral
referral = None
if request.cookies:
ref_code = request.cookies.get(_REFERRAL_COOKIE)
referral = Referral.get_by(code=ref_code)
return referral