From da4e0bf3845b3c2042e9ab533b1f9af1adbe42e4 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Sun, 17 May 2020 10:13:04 +0200
Subject: [PATCH] create /auth/recovery page

---
 app/auth/__init__.py                  |  1 +
 app/auth/templates/auth/recovery.html | 27 +++++++++++
 app/auth/views/recovery.py            | 65 +++++++++++++++++++++++++++
 3 files changed, 93 insertions(+)
 create mode 100644 app/auth/templates/auth/recovery.html
 create mode 100644 app/auth/views/recovery.py

diff --git a/app/auth/__init__.py b/app/auth/__init__.py
index 6a51b706..c41bb56a 100644
--- a/app/auth/__init__.py
+++ b/app/auth/__init__.py
@@ -13,4 +13,5 @@ from .views import (
     mfa,
     fido,
     social,
+    recovery,
 )
diff --git a/app/auth/templates/auth/recovery.html b/app/auth/templates/auth/recovery.html
new file mode 100644
index 00000000..9991c043
--- /dev/null
+++ b/app/auth/templates/auth/recovery.html
@@ -0,0 +1,27 @@
+{% extends "single.html" %}
+
+{% block title %}
+  Recovery Code
+{% endblock %}
+
+
+{% block single_content %}
+  <div class="card">
+    <div class="card-body p-6">
+
+      <form method="post">
+        {{ recovery_form.csrf_token }}
+
+        <div class="font-weight-bold mt-5">Code</div>
+        <div class="small-text">Please enter one of the recovery codes here
+        </div>
+
+        {{ recovery_form.code(class="form-control", autofocus="true") }}
+        {{ render_field_errors(recovery_form.code) }}
+        <button class="btn btn-success mt-2">Submit</button>
+      </form>
+
+    </div>
+  </div>
+
+{% endblock %}
\ No newline at end of file
diff --git a/app/auth/views/recovery.py b/app/auth/views/recovery.py
new file mode 100644
index 00000000..688da290
--- /dev/null
+++ b/app/auth/views/recovery.py
@@ -0,0 +1,65 @@
+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)