Add totp for more account security

config.py

@@ -97,6 +97,7 @@ def load_lines_from_file(filename: str) -> List[str]:
 
 RSS_MOLT_LIMIT = 50
 
+TOTP_ENABLED = getenv_bool('TOTP_ENABLED', False)
 HCAPTCHA_ENABLED = getenv_bool('HCAPTCHA_ENABLED', False)
 REGISTRATION_ENABLED = getenv_bool('REGISTRATION_ENABLED', True)
 

crabber.py

@@ -13,7 +13,9 @@
 from typing import Iterable, Tuple, Union
 import utils
 from werkzeug.middleware.profiler import ProfilerMiddleware
-
+from flask_qrcode import QRcode
+from hashlib import sha256
+from pyotp import TOTP
 
 def create_app():
     app = Flask(__name__, template_folder="./templates")
@@ -27,6 +29,9 @@ def create_app():
     app.config['HCAPTCHA_SECRET_KEY'] = os.getenv('HCAPTCHA_SECRET_KEY')
     app.config['HCAPTCHA_ENABLED'] = config.HCAPTCHA_ENABLED
     app.config['PROFILER_ENABLED'] = os.getenv('PROFILER_ENABLED')
+    app.config['TOTP_SECRET'] = os.getenv('TOTP_SECRET')
+    app.config['TOTP_ISSUER'] = os.getenv('TOTP_ISSUER')
+    app.config['TOTP_ENABLED'] = config.TOTP_ENABLED
 
     register_extensions(app)
     limiter = register_blueprints(app)
@@ -71,7 +76,15 @@ def register_blueprints(app):
 
 
 app, limiter = create_app()
+
+if config.TOTP_ENABLED:
+    secret = sha256()
+    secret.update(app.config['TOTP_SECRET'].encode('utf-8'))
+    global totp; totp = TOTP(secret.hexdigest(),
+        digest=sha256, issuer=app.config['TOTP_ISSUER'])
+
 captcha = hCaptcha(app)
+QRcode(app)
 
 if app.config['PROFILER_ENABLED']:
     app.wsgi_app = ProfilerMiddleware(
@@ -269,11 +282,16 @@ def login():
         if attempted_user is not None:
             if attempted_user.verify_password(password):
                 if not attempted_user.banned:
-                    # Login successful
+                    if app.config['TOTP_ENABLED']:
+                        if attempted_user.totp:
+                            session['totp_user'] = attempted_user.id
+                            session['totp_user_ts'] = \
+                                attempted_user.register_timestamp
+                            return redirect('/login/totp')
                     session['current_user'] = attempted_user.id
                     session['current_user_ts'] = \
                             attempted_user.register_timestamp
-                    return redirect("/")
+                    return redirect("/")              
                 else:
                     return utils.show_error('The account you\'re attempting to'
                                             ' access has been banned.')
@@ -292,6 +310,21 @@ def login():
             login_failed=login_failed
         )
 
+@app.route("/login/totp", methods=('GET', 'POST'))
+def totp_login():
+    if request.method == 'POST':
+        totp_code = request.form.get('totp') or None
+        if (totp_code and totp.verify(totp_code)):
+            session['current_user'] = session['totp_user']
+            session['current_user_ts'] = session['totp_user_ts']
+            del session['totp_user']; del session['totp_user_ts']
+            return redirect('/')
+        else:
+            return utils.show_error('The two factor authentication code given'
+                                    ' was invalid or expired.')
+    if not 'totp_user' in session: return redirect('/login')
+    return render_template('login_2fa.html', 
+        current_page='totp', hide_sidebar=True)
 
 @app.route("/forgotpassword/", methods=('GET', 'POST'))
 def forgot_password():
@@ -547,14 +580,16 @@ def settings():
                 blocks[block] = render_template(
                     f'settings-ajax-{block}.html',
                     current_page='settings',
-                    current_user=utils.get_current_user()
+                    current_user=utils.get_current_user(),
+                    totp=totp
                 )
             return jsonify(blocks)
         else:
             return render_template(
                 'settings.html',
                 current_page='settings',
-                current_user=utils.get_current_user()
+                current_user=utils.get_current_user(),
+                totp=totp
             )
     else:
         return redirect("/login")

models.py

@@ -74,6 +74,9 @@ class Crab(db.Model):
     banned = db.Column(db.Boolean, nullable=False, default=False)
     _password_reset_token = db.Column('password_reset_token', db.String(128))
 
+    # Two Factor Authentication
+    totp = db.Column('two_factor', db.Boolean, nullable=False, default=False)
+
     # Content visibility
     nsfw = db.Column(db.Boolean, nullable=False, default=False)
     show_nsfw = db.Column(db.Boolean, nullable=False, default=False)

pyproject.toml

@@ -22,6 +22,8 @@ beautifulsoup4 = "^4.10.0"
 webpreview = "^1.6.0"
 gunicorn = "^20.1.0"
 gevent = "^21.8.0"
+pyotp = "^2.6.0"
+Flask-QRcode = "^3.0.0"
 
 [tool.poetry.dev-dependencies]
 

templates/login_2fa.html

@@ -0,0 +1,22 @@
+{% extends "base.html" %}
+{% block title %} Two-Factor Authentication {% endblock %}
+{% block heading %}
+<!-- Logo -->
+<img class="logo"
+    src="https://cdn.crabber.net/img/crabber_logo.svg" alt="Crabber Logo" width="43" height="43">
+{% endblock %}
+
+{% block body %}
+    <form class="w-75 m-4" method="POST">
+        <h1 class="mb-5">Two-Factor Authentication</h1>
+        <div class="form-group cool-input">
+            <label for="login-totp">Authentication Code</label>
+            <input type="text" name="totp" class="form-control" id="login-totp" required aria-required>
+        </div>
+        <div class="d-flex align-items-center mt-4">
+            <button type="submit" class="btn btn-primary rounded-pill mr-4">
+                Continue
+            </button>
+        </div>
+    </form>
+{% endblock %}

templates/settings-ajax-body.html

@@ -158,6 +158,21 @@ <h3 class="mb-3">Change password</h3>
     <!-- Section border -->
     <hr class="border-dark mb-3 mt-4">
 
+    <form method="POST" class="p-2">
+        <h3 class="mb-3">Additional Security</h3>
+
+        <div class="custom-control custom-switch">
+            <input type="checkbox" class="custom-control-input" id="toggle-totp"
+                   name="nsfw_mode" {{'checked' if current_user.totp}}
+                   value="{{'true' if current_user.totp else 'false'}}"
+            >
+            <label class="custom-control-label" for="toggle-totp">Two-Factor Authentication</label>
+        </div>
+    </form>
+
+    <!-- Section border -->
+    <hr class="border-dark mb-3 mt-4">
+
     <h2 class="mb-3">Misc</h2>
     <form method="POST" class="p-2">
         <input type="hidden" name="user_action" value="update_general_settings">