From 043ecd4faccd04a4de5e6b3973a9ee089c898bd0 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Sun, 17 May 2020 10:11:38 +0200
Subject: [PATCH] redirect user to recovery codes page after MFA setup. Remove
 all recovery codes when user is no more MFA.

---
 app/dashboard/views/fido_cancel.py | 6 ++++++
 app/dashboard/views/fido_setup.py  | 3 +--
 app/dashboard/views/mfa_cancel.py  | 6 ++++++
 app/dashboard/views/mfa_setup.py   | 3 ++-
 4 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/app/dashboard/views/fido_cancel.py b/app/dashboard/views/fido_cancel.py
index a03cdc50..621aa6f5 100644
--- a/app/dashboard/views/fido_cancel.py
+++ b/app/dashboard/views/fido_cancel.py
@@ -5,6 +5,7 @@ from wtforms import PasswordField, validators
 
 from app.dashboard.base import dashboard_bp
 from app.extensions import db
+from app.models import RecoveryCode
 
 
 class LoginForm(FlaskForm):
@@ -29,6 +30,11 @@ def fido_cancel():
             current_user.fido_sign_count = None
             current_user.fido_credential_id = None
             db.session.commit()
+
+            # user does not have any 2FA enabled left, delete all recovery codes
+            if not current_user.two_factor_authentication_enabled():
+                RecoveryCode.empty(current_user)
+
             flash("We've unlinked your security key.", "success")
             return redirect(url_for("dashboard.index"))
         else:
diff --git a/app/dashboard/views/fido_setup.py b/app/dashboard/views/fido_setup.py
index 5b7fedfe..bf1b00fd 100644
--- a/app/dashboard/views/fido_setup.py
+++ b/app/dashboard/views/fido_setup.py
@@ -68,8 +68,7 @@ def fido_setup():
         db.session.commit()
 
         flash("Security key has been activated", "success")
-
-        return redirect(url_for("dashboard.index"))
+        return redirect(url_for("dashboard.recovery_code_route"))
 
     # Prepare information for key registration process
     fido_uuid = str(uuid.uuid4())
diff --git a/app/dashboard/views/mfa_cancel.py b/app/dashboard/views/mfa_cancel.py
index 6dd940fe..d982b571 100644
--- a/app/dashboard/views/mfa_cancel.py
+++ b/app/dashboard/views/mfa_cancel.py
@@ -6,6 +6,7 @@ from wtforms import StringField, validators
 
 from app.dashboard.base import dashboard_bp
 from app.extensions import db
+from app.models import RecoveryCode
 
 
 class OtpTokenForm(FlaskForm):
@@ -29,6 +30,11 @@ def mfa_cancel():
             current_user.enable_otp = False
             current_user.otp_secret = None
             db.session.commit()
+
+            # user does not have any 2FA enabled left, delete all recovery codes
+            if not current_user.two_factor_authentication_enabled():
+                RecoveryCode.empty(current_user)
+
             flash("MFA is now disabled", "warning")
             return redirect(url_for("dashboard.index"))
         else:
diff --git a/app/dashboard/views/mfa_setup.py b/app/dashboard/views/mfa_setup.py
index 2f0413ab..255e58a1 100644
--- a/app/dashboard/views/mfa_setup.py
+++ b/app/dashboard/views/mfa_setup.py
@@ -36,7 +36,8 @@ def mfa_setup():
             current_user.enable_otp = True
             db.session.commit()
             flash("MFA has been activated", "success")
-            return redirect(url_for("dashboard.index"))
+
+            return redirect(url_for("dashboard.recovery_code_route"))
         else:
             flash("Incorrect token", "warning")