Merge pull request #195 from simple-login/recovery-code

Recovery code

app/auth/__init__.py

@@ -13,4 +13,5 @@ from .views import (
     mfa,
     fido,
     social,
+    recovery,
 )

app/auth/templates/auth/fido.html

@@ -28,11 +28,20 @@
       </div>
 
       {% if enable_otp %}
-        <div class="text-center text-muted mb-6" style="margin-top: 1em;">
+        <hr>
+        <div class="text-muted mt-5" style="margin-top: 1em;">
           Don't have your key with you? <br> <a href="{{ url_for('auth.mfa') }}">Verify by One-Time Password</a>
         </div>
       {% endif %}
 
+      <hr>
+
+      <div class="mt-5">
+        If you have troubles with your authentication app, you can use the recovery code to login. <br>
+        <a href="{{ url_for('auth.recovery_route', next=next_url) }}">Use Recovery Codes</a>
+      </div>
+
+
       <script>
         async function verifyKey() {
           $("#btnVerifyKey").prop('disabled', true);
@@ -65,7 +74,7 @@
       {% if auto_activate %}
         <script>$('document').ready(verifyKey());</script>
       {% endif %}
-        
+
     </div>
   </div>
 

app/auth/templates/auth/mfa.html

@@ -20,7 +20,7 @@
         <input type="hidden" name="form-name" value="create">
 
         <div class="font-weight-bold mt-5">Token</div>
-        <div class="small-text">Please enter the 6-digit number displayed in your MFA application
+        <div class="small-text mb-3">Please enter the 6-digit number displayed in your MFA application
           (Google Authenticator, Authy, MyDigiPassword, etc) here
         </div>
 
@@ -30,11 +30,20 @@
       </form>
 
       {% if enable_fido %}
-        <div class="text-center text-muted mb-6" style="margin-top: 1em;">
+        <hr>
+        <div class="text-muted mt-5" style="margin-top: 1em;">
           Having trouble with your authenticator? <br> <a href="{{ url_for('auth.fido') }}">Verify by your security
           key</a>
         </div>
       {% endif %}
+
+      <hr>
+      <div class="mt-5">
+        If you have troubles with your authentication app, you can use the recovery code to login. <br>
+        <a href="{{ url_for('auth.recovery_route', next=next_url) }}">Use Recovery Codes</a>
+      </div>
+
+
     </div>
   </div>
 

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 %}

app/auth/views/fido.py

@@ -104,4 +104,5 @@ def fido():
         webauthn_assertion_options=webauthn_assertion_options,
         enable_otp=user.enable_otp,
         auto_activate=auto_activate,
+        next_url=next_url,
     )

app/auth/views/login_utils.py

@@ -19,13 +19,13 @@ def after_login(user, next_url):
         # switch between these two 2FA option
         session[MFA_USER_ID] = user.id
         if next_url:
-            return redirect(url_for("auth.fido", next_url=next_url))
+            return redirect(url_for("auth.fido", next=next_url))
         else:
             return redirect(url_for("auth.fido"))
     elif user.enable_otp:
         session[MFA_USER_ID] = user.id
         if next_url:
-            return redirect(url_for("auth.mfa", next_url=next_url))
+            return redirect(url_for("auth.mfa", next=next_url))
         else:
             return redirect(url_for("auth.mfa"))
     else:

app/auth/views/mfa.py

@@ -59,4 +59,5 @@ def mfa():
         "auth/mfa.html",
         otp_token_form=otp_token_form,
         enable_fido=(user.fido_enabled()),
+        next_url=next_url,
     )

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)

app/dashboard/__init__.py

@@ -20,4 +20,5 @@ from .views import (
     mailbox_detail,
     refused_email,
     referral,
+    recovery_code,
 )

app/dashboard/templates/dashboard/recovery_code.html

@@ -0,0 +1,40 @@
+{% extends 'default.html' %}
+{% set active_page = "setting" %}
+{% block title %}
+  Recovery Codes
+{% endblock %}
+
+{% block default_content %}
+  <div class="card">
+    <div class="card-body">
+      <h1 class="h3">Recovery codes</h1>
+      <p>
+        If you lose access to your authentication device, you can use one of these backup codes to login to your
+        account. Each code may be used only once. Make a copy of these codes, and store it somewhere safe.
+      </p>
+
+      <ul>
+        {% for recovery_code in recovery_codes %}
+          {% if recovery_code.used %}
+            <li>
+              <span style="text-decoration: line-through">{{ recovery_code.code }}</span>.
+              Used {{ recovery_code.used_at | dt }}.
+            </li>
+          {% else %}
+            <li>
+              {{ recovery_code.code }}
+            </li>
+          {% endif %}
+        {% endfor %}
+      </ul>
+
+      <form method="post" class="mt-6">
+        <input type="submit" class="btn btn-outline-primary" value="Generate New Codes">
+      </form>
+      <div class="small-text">
+        Beware: Generating new codes invalidates all the previous ones, make sure to write down the new ones!
+      </div>
+
+    </div>
+  </div>
+{% endblock %}

app/dashboard/templates/dashboard/setting.html

@@ -98,6 +98,7 @@
             <a href="{{ url_for('dashboard.fido_setup') }}" class="btn btn-outline-primary">Setup WebAuthn</a>
           {% else %}
             <a href="{{ url_for('dashboard.fido_cancel') }}" class="btn btn-outline-danger">Disable WebAuthn</a>
+            <a href="{{ url_for('dashboard.recovery_code_route') }}" class="btn btn-outline-secondary">Recovery Codes</a>
           {% endif %}
         </div>
       </div>
@@ -114,6 +115,7 @@
           <a href="{{ url_for('dashboard.mfa_setup') }}" class="btn btn-outline-primary">Setup TOTP</a>
         {% else %}
           <a href="{{ url_for('dashboard.mfa_cancel') }}" class="btn btn-outline-danger">Disable TOTP</a>
+          <a href="{{ url_for('dashboard.recovery_code_route') }}" class="btn btn-outline-secondary">Recovery Codes</a>
         {% endif %}
       </div>
     </div>

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:

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())

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:

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")
 

app/dashboard/views/recovery_code.py

@@ -0,0 +1,30 @@
+from flask import render_template, flash, redirect, url_for, request
+from flask_login import login_required, current_user
+
+from app.dashboard.base import dashboard_bp
+from app.log import LOG
+from app.models import RecoveryCode
+
+
+@dashboard_bp.route("/recovery_code", methods=["GET", "POST"])
+@login_required
+def recovery_code_route():
+    if not current_user.two_factor_authentication_enabled():
+        flash("you need to enable either TOTP or WebAuthn", "warning")
+        return redirect(url_for("dashboard.index"))
+
+    recovery_codes = RecoveryCode.query.filter_by(user_id=current_user.id).all()
+    if request.method == "GET" and not recovery_codes:
+        # user arrives at this page for the first time
+        LOG.d("%s has no recovery keys, generate", current_user)
+        RecoveryCode.generate(current_user)
+        recovery_codes = RecoveryCode.query.filter_by(user_id=current_user.id).all()
+
+    if request.method == "POST":
+        RecoveryCode.generate(current_user)
+        flash("New recovery codes generated", "success")
+        return redirect(url_for("dashboard.recovery_code_route"))
+
+    return render_template(
+        "dashboard/recovery_code.html", recovery_codes=recovery_codes
+    )

app/models.py

@@ -164,6 +164,9 @@ class User(db.Model, ModelMixin, UserMixin):
             return True
         return False
 
+    def two_factor_authentication_enabled(self) -> bool:
+        return self.enable_otp or self.fido_enabled()
+
     # some users could have lifetime premium
     lifetime = db.Column(db.Boolean, default=False, nullable=False, server_default="0")
 
@@ -1344,3 +1347,43 @@ class AliasMailbox(db.Model, ModelMixin):
     mailbox_id = db.Column(
         db.ForeignKey(Mailbox.id, ondelete="cascade"), nullable=False
     )
+
+
+_NB_RECOVERY_CODE = 8
+_RECOVERY_CODE_LENGTH = 8
+
+
+class RecoveryCode(db.Model, ModelMixin):
+    """allow user to login in case you lose any of your authenticators"""
+
+    __table_args__ = (db.UniqueConstraint("user_id", "code", name="uq_recovery_code"),)
+
+    user_id = db.Column(db.ForeignKey(User.id, ondelete="cascade"), nullable=False)
+    code = db.Column(db.String(16), nullable=False)
+    used = db.Column(db.Boolean, nullable=False, default=False)
+    used_at = db.Column(ArrowType, nullable=True, default=None)
+
+    user = db.relationship(User)
+
+    @classmethod
+    def generate(cls, user):
+        """generate recovery codes for user"""
+        # delete all existing codes
+        cls.query.filter_by(user_id=user.id).delete()
+        db.session.flush()
+
+        nb_code = 0
+        while nb_code < _NB_RECOVERY_CODE:
+            code = random_string(_RECOVERY_CODE_LENGTH)
+            if not cls.get_by(user_id=user.id, code=code):
+                cls.create(user_id=user.id, code=code)
+                nb_code += 1
+
+        LOG.d("Create recovery codes for %s", user)
+        db.session.commit()
+
+    @classmethod
+    def empty(cls, user):
+        """Delete all recovery codes for user"""
+        cls.query.filter_by(user_id=user.id).delete()
+        db.session.commit()

migrations/versions/2020_051710_c31cdf879ee3_.py

@@ -0,0 +1,40 @@
+"""empty message
+
+Revision ID: c31cdf879ee3
+Revises: 5cad8fa84386
+Create Date: 2020-05-17 10:34:23.492008
+
+"""
+import sqlalchemy_utils
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'c31cdf879ee3'
+down_revision = '5cad8fa84386'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('recovery_code',
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('created_at', sqlalchemy_utils.types.arrow.ArrowType(), nullable=False),
+    sa.Column('updated_at', sqlalchemy_utils.types.arrow.ArrowType(), nullable=True),
+    sa.Column('user_id', sa.Integer(), nullable=False),
+    sa.Column('code', sa.String(length=16), nullable=False),
+    sa.Column('used', sa.Boolean(), nullable=False),
+    sa.Column('used_at', sqlalchemy_utils.types.arrow.ArrowType(), nullable=True),
+    sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='cascade'),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('user_id', 'code', name='uq_recovery_code')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('recovery_code')
+    # ### end Alembic commands ###