mirror of
https://github.com/simple-login/app.git
synced 2024-11-10 21:27:10 +01:00
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
|
import arrow
|
||
|
|
import pyotp
|
||
|
|
from flask import request, render_template, redirect, url_for, flash, session
|
||
|
|
from flask_login import login_user
|
||
|
|
from flask_wtf import FlaskForm
|
||
|
|
from wtforms import StringField, validators
|
||
|
|
|
||
|
|
from app.auth.base import auth_bp
|
||
|
|
from app.config import MFA_USER_ID
|
||
|
|
from app.extensions import db
|
||
|
|
from app.log import LOG
|
||
|
|
from app.models import User, RecoveryCode
|
||
|
|
|
||
|
|
|
||
|
|
class RecoveryForm(FlaskForm):
|
||
|
|
code = StringField("Code", validators=[validators.DataRequired()])
|
||
|
|
|
||
|
|
|
||
|
|
@auth_bp.route("/recovery", methods=["GET", "POST"])
|
||
|
|
def recovery_route():
|
||
|
|
# passed from login page
|
||
|
|
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")
|
||
|
|
return redirect(url_for("auth.login"))
|
||
|
|
|
||
|
|
user = User.get(user_id)
|
||
|
|
|
||
|
|
if not user.two_factor_authentication_enabled():
|
||
|
|
flash("Only user with MFA enabled should go to this page", "warning")
|
||
|
|
return redirect(url_for("auth.login"))
|
||
|
|
|
||
|
|
recovery_form = RecoveryForm()
|
||
|
|
next_url = request.args.get("next")
|
||
|
|
|
||
|
|
if recovery_form.validate_on_submit():
|
||
|
|
code = recovery_form.code.data
|
||
|
|
recovery_code = RecoveryCode.get_by(user_id=user.id, code=code)
|
||
|
|
|
||
|
|
if recovery_code:
|
||
|
|
if recovery_code.used:
|
||
|
|
flash("Code already used", "error")
|
||
|
|
else:
|
||
|
|
del session[MFA_USER_ID]
|
||
|
|
|
||
|
|
login_user(user)
|
||
|
|
flash(f"Welcome back {user.name}!", "success")
|
||
|
|
|
||
|
|
recovery_code.used = True
|
||
|
|
recovery_code.used_at = arrow.now()
|
||
|
|
db.session.commit()
|
||
|
|
|
||
|
|
# 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"))
|
||
|
|
else:
|
||
|
|
flash("Incorrect code", "error")
|
||
|
|
|
||
|
|
return render_template("auth/recovery.html", recovery_form=recovery_form)
|