diff --git a/app/auth/routes.py b/app/auth/routes.py index 675c9bf..d1327d4 100644 --- a/app/auth/routes.py +++ b/app/auth/routes.py @@ -234,6 +234,64 @@ def disable_2fa(): return redirect(url_for('main.profile')) +@bp.route('/enable-2fa') +@login_required +def enable_2fa(): + """ + Enable 2FA for existing logged-in users. + + Security: Requires active login session. + """ + if current_user.is_2fa_enabled: + flash('Two-factor authentication is already enabled for your account.', 'info') + return redirect(url_for('main.profile')) + + # Generate new TOTP secret if not exists + if not current_user.totp_secret: + current_user.generate_totp_secret() + db.session.commit() + + # Generate QR code for Google Authenticator + qr_code = current_user.generate_qr_code() + + return render_template('auth/enable_2fa.html', + qr_code=qr_code, + username=current_user.username, + title='Enable Two-Factor Authentication') + + +@bp.route('/verify-enable-2fa', methods=['POST']) +@login_required +def verify_enable_2fa(): + """ + Verify TOTP token and enable 2FA for existing user. + + Security: Requires active login session and valid TOTP token. + """ + if current_user.is_2fa_enabled: + flash('Two-factor authentication is already enabled.', 'info') + return redirect(url_for('main.profile')) + + token = request.form.get('token', '').strip() + + if not token: + flash('Please enter the authentication code.', 'error') + return redirect(url_for('auth.enable_2fa')) + + if current_user.verify_totp(token): + # Enable 2FA for the user + current_user.enable_2fa() + + logger.info(f'2FA enabled for user: {current_user.username}') + flash('Two-factor authentication has been successfully enabled!', 'success') + + return redirect(url_for('main.profile')) + else: + logger.warning(f'Failed 2FA enable verification for user: {current_user.username}') + flash('Invalid authentication code. Please try again.', 'error') + return redirect(url_for('auth.enable_2fa')) + + def track_login_location(user): """ Track user login location and handle suspicious login detection. diff --git a/app/models.py b/app/models.py index 19a3e84..7141fae 100644 --- a/app/models.py +++ b/app/models.py @@ -94,13 +94,12 @@ class User(UserMixin, db.Model): totp = pyotp.TOTP(self.totp_secret) # Verify token with a 1-period window (30 seconds before/after) return totp.verify(token, valid_window=1) - def generate_qr_code(self, issuer_name='Flask-2FA-App'): """ Generate a QR code for the TOTP URI. Returns: - str: Base64-encoded PNG image of the QR code + str: HTML img tag with base64-encoded PNG image of the QR code """ uri = self.generate_totp_uri(issuer_name) @@ -122,7 +121,8 @@ class User(UserMixin, db.Model): img.save(img_buffer, format='PNG') img_buffer.seek(0) - return base64.b64encode(img_buffer.getvalue()).decode() + base64_img = base64.b64encode(img_buffer.getvalue()).decode() + return f'QR Code for 2FA Setup' def enable_2fa(self): """Enable two-factor authentication for the user.""" diff --git a/app/templates/auth/enable_2fa.html b/app/templates/auth/enable_2fa.html new file mode 100644 index 0000000..4eb4a66 --- /dev/null +++ b/app/templates/auth/enable_2fa.html @@ -0,0 +1,157 @@ +{% extends "base.html" %} + +{% block content %} +
+
+
+
+

Enable Two-Factor Authentication

+

Secure your account with an additional layer of protection

+
+
+ {% with messages = get_flashed_messages(with_categories=true) %} + {% if messages %} + {% for category, message in messages %} +
+ {{ message }} + +
+ {% endfor %} + {% endif %} + {% endwith %} + +
+
+
Step 1: Scan QR Code
+
+
+ {{ qr_code|safe }} +
+
+ +
+
Instructions:
+
    +
  1. Install an authenticator app like Google Authenticator or Authy
    2. +
  2. Open the app and scan this QR code
    3. +
  3. Enter the 6-digit code from your app below
    4. +
+
+
+ +
+
Step 2: Verify Setup
+ +
+ + +
+ + +
Enter the 6-digit code from your authenticator app
+
+ +
+ + + Cancel + +
+
+
+
+ +
+ +
+
+
Security Benefits
+
+
+
+ +
Enhanced Security
+ Protect against password theft +
+
+
+
+ +
Mobile Protection
+ Secure codes on your device +
+
+
+
+ +
Location Tracking
+ Monitor suspicious logins +
+
+
+
+
+
+
+
+
+ + + + +{% endblock %} diff --git a/app/templates/profile.html b/app/templates/profile.html index c70a97d..fafa9f3 100644 --- a/app/templates/profile.html +++ b/app/templates/profile.html @@ -68,22 +68,20 @@ Enabled
Your account is protected with 2FA - -
- {{ csrf_token() }} +
- {% else %} -
+ {% else %}
Setup Required
Complete 2FA setup for better security
- + Setup 2FA {% endif %}