Allow user to disable mfa for browser for 30 days

app/auth/templates/auth/mfa.html

@@ -26,6 +26,12 @@
 
         {{ otp_token_form.token(class="form-control", autofocus="true") }}
         {{ render_field_errors(otp_token_form.token) }}
+        <div class="form-check">
+          {{ otp_token_form.remember(class="form-check-input", id="remember") }}
+          <label class="form-check-label" for="remember">
+            {{ otp_token_form.remember.description }}
+          </label>
+        </div>
         <button class="btn btn-success mt-2">Validate</button>
       </form>
 

app/auth/views/logout.py

@@ -1,4 +1,4 @@
-from flask import redirect, url_for, flash
+from flask import redirect, url_for, flash, make_response
 from flask_login import logout_user
 
 from app.auth.base import auth_bp
@@ -8,4 +8,9 @@ from app.auth.base import auth_bp
 def logout():
     logout_user()
     flash("You are logged out", "success")
-    return redirect(url_for("auth.login"))
+    response = make_response(redirect(url_for("auth.login")))
+    response.delete_cookie("slapp")
+    response.delete_cookie("mfa")
+    response.delete_cookie("dark-mode")
+
+    return response

app/auth/views/mfa.py

@@ -1,17 +1,29 @@
 import pyotp
-from flask import request, render_template, redirect, url_for, flash, session
+from flask import (
+    render_template,
+    redirect,
+    url_for,
+    flash,
+    session,
+    make_response,
+    request,
+)
 from flask_login import login_user
 from flask_wtf import FlaskForm
-from wtforms import StringField, validators
+from wtforms import BooleanField, StringField, validators
 
 from app.auth.base import auth_bp
-from app.config import MFA_USER_ID
+from app.config import MFA_USER_ID, URL
+from app.extensions import db
 from app.log import LOG
-from app.models import User
+from app.models import User, MfaBrowser
 
 
 class OtpTokenForm(FlaskForm):
     token = StringField("Token", validators=[validators.DataRequired()])
+    remember = BooleanField(
+        "attr", default=False, description="Remember this browser for 30 days"
+    )
 
 
 @auth_bp.route("/mfa", methods=["GET", "POST"])
@@ -33,6 +45,16 @@ def mfa():
     otp_token_form = OtpTokenForm()
     next_url = request.args.get("next")
 
+    if request.cookies.get("mfa"):
+        browser = MfaBrowser.get_by(token=request.cookies.get("mfa"))
+        if browser and not browser.is_expired():
+            login_user(user)
+            flash(f"Welcome back {user.name}!", "success")
+            # Redirect user to correct page
+            return redirect(next_url or url_for("dashboard.index"))
+
+        MfaBrowser.delete(browser.token)
+
     if otp_token_form.validate_on_submit():
         totp = pyotp.TOTP(user.otp_secret)
 
@@ -42,15 +64,24 @@ def mfa():
             del session[MFA_USER_ID]
 
             login_user(user)
-            flash(f"Welcome back {user.name}!")
+            flash(f"Welcome back {user.name}!", "success")
 
-            # 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"))
+            # Redirect user to correct page
+            response = make_response(redirect(next_url or url_for("dashboard.index")))
+
+            if otp_token_form.remember.data:
+                browser = MfaBrowser.create_new(user=user)
+                db.session.commit()
+                response.set_cookie(
+                    "mfa",
+                    value=browser.token,
+                    expires=browser.expires.datetime,
+                    secure=True if URL.startswith("https") else False,
+                    httponly=True,
+                    samesite="Lax",
+                )
+
+            return response
 
         else:
             flash("Incorrect token", "warning")

app/models.py

@@ -76,9 +76,9 @@ class ModelMixin(object):
     def save(self):
         db.session.add(self)
 
-    @classmethod
-    def delete(cls, obj_id):
-        cls.query.filter(cls.id == obj_id).delete()
+        @classmethod
+        def delete(cls, obj_id):
+            cls.query.filter(cls.id == obj_id).delete()
 
     def __repr__(self):
         values = ", ".join(
@@ -508,6 +508,38 @@ def generate_oauth_client_id(client_name) -> str:
     return generate_oauth_client_id(client_name)
 
 
+class MfaBrowser(db.Model, ModelMixin):
+    user_id = db.Column(db.ForeignKey(User.id, ondelete="cascade"), nullable=False)
+    token = db.Column(db.String(64), default=False, nullable=False)
+    expires = db.Column(ArrowType, default=False, nullable=False)
+
+    user = db.relationship(User)
+
+    @classmethod
+    def create_new(cls, user, token_length=64) -> "MfaBrowser":
+        return MfaBrowser.create(
+            user_id=user.id,
+            token=random_string(token_length),
+            expires=arrow.now().shift(days=30),
+        )
+
+    @classmethod
+    def delete(cls, token):
+        cls.query.filter(cls.token == token).delete()
+        db.session.commit()
+
+    @classmethod
+    def delete_expired(cls):
+        cls.query.filter(cls.expires < arrow.now()).delete()
+        db.session.commit()
+
+    def is_expired(self):
+        return self.expires < arrow.now()
+
+    def reset_expire(self):
+        self.expires = arrow.now().shift(days=30)
+
+
 class Client(db.Model, ModelMixin):
     oauth_client_id = db.Column(db.String(128), unique=True, nullable=False)
     oauth_client_secret = db.Column(db.String(128), nullable=False)

migrations/versions/2020_052216_95599239860a_.py

@@ -0,0 +1,38 @@
+"""empty message
+
+Revision ID: 95599239860a
+Revises: ce15cf3467b4
+Create Date: 2020-05-22 16:14:33.704035
+
+"""
+import sqlalchemy_utils
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '95599239860a'
+down_revision = 'ce15cf3467b4'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('mfa_browser',
+    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('token', sa.String(length=64), nullable=False),
+    sa.Column('expires', sqlalchemy_utils.types.arrow.ArrowType(), nullable=False),
+    sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='cascade'),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('mfa_browser')
+    # ### end Alembic commands ###