Create MFA Setup page if user enable MFA

app/config.py

@@ -113,3 +113,4 @@ AVATAR_URL_EXPIRATION = 3600 * 24 * 7  # 1h*24h/d*7d=1week
 
 # session key
 HIGHLIGHT_GEN_EMAIL_ID = "highlight_gen_email_id"
+MFA_USER_ID = "mfa_user_id"

app/dashboard/__init__.py

@@ -9,4 +9,5 @@ from .views import (
     api_key,
     custom_domain,
     alias_contact_manager,
+    mfa_setup,
 )

app/dashboard/templates/dashboard/mfa_setup.html

@@ -0,0 +1,51 @@
+{% extends 'default.html' %}
+
+{% block title %}
+  MFA Setup
+{% endblock %}
+
+{% block head %}
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/qrious/4.0.2/qrious.min.js"></script>
+{% endblock %}
+
+{% block default_content %}
+  <div class="bg-white p-6" style="max-width: 60em; margin: auto">
+    <h1 class="h2">Two Factor Authentication</h1>
+    <p>Please open a TOTP application (Google Authenticator, Authy, etc)
+      on your smartphone and scan the following QR Code:
+    </p>
+
+    <canvas id="qr"></canvas>
+
+    <script>
+      (function () {
+        var qr = new QRious({
+          element: document.getElementById('qr'),
+          value: '{{otp_uri}}'
+        });
+      })();
+    </script>
+
+    <div class="mt-3 mb-2">
+      Or you can use the manual entry with the following key:
+    </div>
+
+    <div class="mb-3 p-3" style="background-color: #eee">
+      {{ current_user.otp_secret }}
+    </div>
+
+
+    <form method="post">
+      {{ otp_token_form.csrf_token }}
+
+      <div class="font-weight-bold mt-5">Token</div>
+      <div class="small-text">Please enter the 6-digit number displayed on your phone.</div>
+
+      {{ otp_token_form.token(class="form-control", placeholder="") }}
+      {{ render_field_errors(otp_token_form.token) }}
+      <button class="btn btn-lg btn-success mt-2">Validate</button>
+    </form>
+
+
+  </div>
+{% endblock %}

app/dashboard/views/mfa_setup.py

@@ -0,0 +1,49 @@
+import pyotp
+from flask import render_template, flash, redirect, url_for
+from flask_login import login_required, current_user
+from flask_wtf import FlaskForm
+from wtforms import StringField, validators
+
+from app.dashboard.base import dashboard_bp
+from app.extensions import db
+from app.log import LOG
+
+
+class OtpTokenForm(FlaskForm):
+    token = StringField("Token", validators=[validators.DataRequired()])
+
+
+@dashboard_bp.route("/mfa_setup", methods=["GET", "POST"])
+@login_required
+def mfa_setup():
+    if current_user.enable_otp:
+        flash("you have already enabled MFA", "warning")
+        return redirect(url_for("dashboard.index"))
+
+    otp_token_form = OtpTokenForm()
+
+    if not current_user.otp_secret:
+        LOG.d("Generate otp_secret for user %s", current_user)
+        current_user.otp_secret = pyotp.random_base32()
+        db.session.commit()
+
+    totp = pyotp.TOTP(current_user.otp_secret)
+
+    if otp_token_form.validate_on_submit():
+        token = otp_token_form.token.data
+
+        if totp.verify(token):
+            current_user.enable_otp = True
+            db.session.commit()
+            flash("2FA has been activated", "success")
+            return redirect(url_for("dashboard.index"))
+        else:
+            flash("Incorrect token", "warning")
+
+    otp_uri = pyotp.totp.TOTP(current_user.otp_secret).provisioning_uri(
+        name=current_user.email, issuer_name="SimpleLogin"
+    )
+
+    return render_template(
+        "dashboard/mfa_setup.html", otp_token_form=otp_token_form, otp_uri=otp_uri
+    )