Update login workflow to use 2FA

hiboo/account/forms.py

@@ -11,6 +11,11 @@ class LoginForm(flask_wtf.FlaskForm):
     submit = fields.SubmitField(_('Sign in'))
 
 
+class TotpForm(flask_wtf.FlaskForm):
+    totp = fields.PasswordField(_('Time-based One-Time Password'), [validators.DataRequired()])
+    submit = fields.SubmitField(_('Validate'))
+
+
 class SignupForm(flask_wtf.FlaskForm):
     username = fields.StringField(_('Username'), [
         validators.DataRequired(),

hiboo/account/login.py

@@ -14,6 +14,9 @@ def signin():
     if form.validate_on_submit():
         user = models.User.login(form.username.data, form.password.data)
         if user:
+            if "totp" in user.auths:
+                session["username"] = user.username
+                return flask.redirect(flask.url_for(".totp_verify"))
             flask_login.login_user(user)
             if form.remember_me.data == True:
                 session.permanent = True
@@ -24,6 +27,23 @@ def signin():
         action=utils.url_for(".signin"))
 
 
+@blueprint.route("/totp/verify", methods=["GET", "POST"])
+def totp_verify():
+    form = forms.TotpForm()
+    if "username" not in session:
+        return flask.redirect(flask.url_for(".signin"))
+    if form.validate_on_submit():
+        user = models.User.query.filter_by(username=session["username"]).first()
+        if user and user.auths["totp"].check_totp(form.totp.data):
+            flask_login.login_user(user)
+            session.pop("username")
+            return flask.redirect(utils.url_or_intent(".home"))
+        else:
+            flask.flash(_("Wrong password"), "danger")
+    return flask.render_template("account_totp_verify.html", form=form,
+        action=utils.url_for(".totp_verify"))
+
+
 @blueprint.route("/signout")
 @security.authentication_required()
 def signout():

hiboo/account/templates/account_totp_verify.html

+{% extends "base.html" %}
+
+{% block title %}{% trans %}Time-based One-Time Password (TOTP) verify{% endtrans %}{% endblock %}
+{% block subtitle %}{% trans %}to access your account{% endtrans %}{% endblock %}
+
+{% block content %}
+
+{{ macros.form(form) }}
+
+{% endblock %}
+
+{% block actions %}
+<a href="{{ utils.url_for(".signup") }}" class="btn btn-success">{% trans %}Sign up{% endtrans %}</a>
+{% endblock %}