from time import time 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 from app.models import Referral def after_login(user, next_url, login_from_proton: bool = False): """ Redirect to the correct page after login. If the user is logged in with Proton, do not look at fido nor otp If user enables MFA: redirect user to MFA page Otherwise redirect to dashboard page if no next_url """ if not login_from_proton: if user.fido_enabled(): # Use the same session for FIDO so that we can easily # switch between these two 2FA option session[MFA_USER_ID] = user.id if next_url: return redirect(url_for("auth.fido", next=next_url)) else: return redirect(url_for("auth.fido")) elif user.enable_otp: session[MFA_USER_ID] = user.id if next_url: return redirect(url_for("auth.mfa", next=next_url)) else: return redirect(url_for("auth.mfa")) LOG.d("log user %s in", user) login_user(user) session["sudo_time"] = int(time()) # User comes to login page from another page if next_url: LOG.d("redirect user to %s", next_url) return redirect(next_url) else: LOG.d("redirect user to dashboard") return redirect(url_for("dashboard.index")) # 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) if not referral: if "slref" in session: ref_code = session["slref"] referral = Referral.get_by(code=ref_code) if referral: LOG.d("referral found %s", referral) return referral