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

105 lines
3.1 KiB
Python
Raw Normal View History

import pyotp
from flask import (
render_template,
redirect,
url_for,
flash,
session,
make_response,
request,
2020-05-24 16:14:48 +02:00
g,
)
from flask_login import login_user
from flask_wtf import FlaskForm
from wtforms import BooleanField, StringField, validators
from app.auth.base import auth_bp
from app.config import MFA_USER_ID, URL
from app.db import Session
from app.extensions import limiter
from app.models import User, MfaBrowser
class OtpTokenForm(FlaskForm):
token = StringField("Token", validators=[validators.DataRequired()])
remember = BooleanField(
"attr", default=False, description="Remember this browser for 30 days"
)
@auth_bp.route("/mfa", methods=["GET", "POST"])
2020-05-24 16:14:48 +02:00
@limiter.limit(
"10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
)
def mfa():
# passed from login page
2020-01-03 23:42:35 +01:00
user_id = session.get(MFA_USER_ID)
# user access this page directly without passing by login page
if not user_id:
flash("Unknown error, redirect back to main page", "warning")
2020-01-03 23:50:34 +01:00
return redirect(url_for("auth.login"))
2020-01-03 23:42:35 +01:00
user = User.get(user_id)
2020-01-03 23:42:35 +01:00
if not (user and user.enable_otp):
flash("Only user with MFA enabled should go to this page", "warning")
2020-01-03 23:50:34 +01:00
return redirect(url_for("auth.login"))
otp_token_form = OtpTokenForm()
next_url = request.args.get("next")
if request.cookies.get("mfa"):
browser = MfaBrowser.get_by(token=request.cookies.get("mfa"))
if browser and not browser.is_expired() and browser.user_id == user.id:
login_user(user)
flash(f"Welcome back!", "success")
# Redirect user to correct page
return redirect(next_url or url_for("dashboard.index"))
2020-05-24 16:14:48 +02:00
else:
# Trigger rate limiter
g.deduct_limit = True
if otp_token_form.validate_on_submit():
totp = pyotp.TOTP(user.otp_secret)
token = otp_token_form.token.data.replace(" ", "")
if totp.verify(token) and user.last_otp != token:
2019-12-27 17:36:35 +01:00
del session[MFA_USER_ID]
user.last_otp = token
Session.commit()
2019-12-27 17:36:35 +01:00
login_user(user)
flash(f"Welcome back!", "success")
# Redirect user to correct page
response = make_response(redirect(next_url or url_for("dashboard.index")))
if otp_token_form.remember.data:
browser = MfaBrowser.create_new(user=user)
Session.commit()
response.set_cookie(
"mfa",
value=browser.token,
expires=browser.expires.datetime,
secure=True if URL.startswith("https") else False,
httponly=True,
samesite="Lax",
)
return response
2019-12-27 17:36:35 +01:00
else:
flash("Incorrect token", "warning")
2020-05-24 16:14:48 +02:00
# Trigger rate limiter
g.deduct_limit = True
otp_token_form.token.data = None
2020-05-07 11:53:28 +02:00
return render_template(
"auth/mfa.html",
otp_token_form=otp_token_form,
2020-05-07 14:32:52 +02:00
enable_fido=(user.fido_enabled()),
2020-05-17 10:35:11 +02:00
next_url=next_url,
2020-05-07 11:53:28 +02:00
)