From 0964a2177a93b300e8298f2d46fc837fd7587427 Mon Sep 17 00:00:00 2001 From: Jens Luedicke Date: Sun, 29 Jun 2025 15:41:55 +0200 Subject: [PATCH] Add 2FA authentification. --- app.py | 122 ++++++++++++++++- migrate_db.py | 66 +++++++++- models.py | 32 +++++ requirements.txt | 4 +- templates/profile.html | 109 ++++++++++++++++ templates/setup_2fa.html | 266 ++++++++++++++++++++++++++++++++++++++ templates/verify_2fa.html | 182 ++++++++++++++++++++++++++ 7 files changed, 768 insertions(+), 13 deletions(-) create mode 100644 templates/setup_2fa.html create mode 100644 templates/verify_2fa.html diff --git a/app.py b/app.py index 3811234..d5245b4 100644 --- a/app.py +++ b/app.py @@ -191,14 +191,20 @@ def login(): if user.is_blocked: flash('Your account has been disabled. Please contact an administrator.', 'error') return render_template('login.html') - - # Continue with normal login process - session['user_id'] = user.id - session['username'] = user.username - session['is_admin'] = user.is_admin - flash('Login successful!', 'success') - return redirect(url_for('home')) + # Check if 2FA is enabled + if user.two_factor_enabled: + # Store user ID for 2FA verification + session['2fa_user_id'] = user.id + return redirect(url_for('verify_2fa')) + else: + # Continue with normal login process + session['user_id'] = user.id + session['username'] = user.username + session['is_admin'] = user.is_admin + + flash('Login successful!', 'success') + return redirect(url_for('home')) flash('Invalid username or password', 'error') @@ -552,6 +558,108 @@ def profile(): return render_template('profile.html', title='My Profile', user=user) +@app.route('/2fa/setup', methods=['GET', 'POST']) +@login_required +def setup_2fa(): + if request.method == 'POST': + # Verify the TOTP code before enabling 2FA + totp_code = request.form.get('totp_code') + + if not totp_code: + flash('Please enter the verification code from your authenticator app.', 'error') + return redirect(url_for('setup_2fa')) + + try: + if g.user.verify_2fa_token(totp_code, allow_setup=True): + g.user.two_factor_enabled = True + db.session.commit() + flash('Two-factor authentication has been successfully enabled!', 'success') + return redirect(url_for('profile')) + else: + flash('Invalid verification code. Please make sure your device time is synchronized and try again.', 'error') + return redirect(url_for('setup_2fa')) + except Exception as e: + logger.error(f"2FA setup error: {str(e)}") + flash('An error occurred during 2FA setup. Please try again.', 'error') + return redirect(url_for('setup_2fa')) + + # GET request - show setup page + if g.user.two_factor_enabled: + flash('Two-factor authentication is already enabled.', 'info') + return redirect(url_for('profile')) + + # Generate secret if not exists + if not g.user.two_factor_secret: + g.user.generate_2fa_secret() + db.session.commit() + + # Generate QR code + import qrcode + import io + import base64 + + qr_uri = g.user.get_2fa_uri() + qr = qrcode.QRCode(version=1, box_size=10, border=5) + qr.add_data(qr_uri) + qr.make(fit=True) + + # Create QR code image + qr_img = qr.make_image(fill_color="black", back_color="white") + img_buffer = io.BytesIO() + qr_img.save(img_buffer, format='PNG') + img_buffer.seek(0) + qr_code_b64 = base64.b64encode(img_buffer.getvalue()).decode() + + return render_template('setup_2fa.html', + title='Setup Two-Factor Authentication', + secret=g.user.two_factor_secret, + qr_code=qr_code_b64) + +@app.route('/2fa/disable', methods=['POST']) +@login_required +def disable_2fa(): + password = request.form.get('password') + + if not password or not g.user.check_password(password): + flash('Please enter your correct password to disable 2FA.', 'error') + return redirect(url_for('profile')) + + g.user.two_factor_enabled = False + g.user.two_factor_secret = None + db.session.commit() + + flash('Two-factor authentication has been disabled.', 'success') + return redirect(url_for('profile')) + +@app.route('/2fa/verify', methods=['GET', 'POST']) +def verify_2fa(): + # Check if user is in 2FA verification state + user_id = session.get('2fa_user_id') + if not user_id: + return redirect(url_for('login')) + + user = User.query.get(user_id) + if not user or not user.two_factor_enabled: + session.pop('2fa_user_id', None) + return redirect(url_for('login')) + + if request.method == 'POST': + totp_code = request.form.get('totp_code') + + if user.verify_2fa_token(totp_code): + # Complete login process + session.pop('2fa_user_id', None) + session['user_id'] = user.id + session['username'] = user.username + session['is_admin'] = user.is_admin + + flash('Login successful!', 'success') + return redirect(url_for('home')) + else: + flash('Invalid verification code. Please try again.', 'error') + + return render_template('verify_2fa.html', title='Two-Factor Authentication') + @app.route('/about') @login_required def about(): diff --git a/migrate_db.py b/migrate_db.py index 4734cb3..703f85a 100644 --- a/migrate_db.py +++ b/migrate_db.py @@ -1,7 +1,7 @@ from app import app, db import sqlite3 import os -from models import User, TimeEntry, WorkConfig, SystemSettings +from models import User, TimeEntry, WorkConfig, SystemSettings, Team, Role from werkzeug.security import generate_password_hash from datetime import datetime @@ -106,6 +106,38 @@ def migrate_database(): if 'is_blocked' not in user_columns: print("Adding is_blocked column to user table...") cursor.execute("ALTER TABLE user ADD COLUMN is_blocked BOOLEAN DEFAULT 0") + + # Add role column to user table if it doesn't exist + if 'role' not in user_columns: + print("Adding role column to user table...") + cursor.execute("ALTER TABLE user ADD COLUMN role VARCHAR(50) DEFAULT 'Team Member'") + + # Add team_id column to user table if it doesn't exist + if 'team_id' not in user_columns: + print("Adding team_id column to user table...") + cursor.execute("ALTER TABLE user ADD COLUMN team_id INTEGER") + + # Add 2FA columns to user table if they don't exist + if 'two_factor_enabled' not in user_columns: + print("Adding two_factor_enabled column to user table...") + cursor.execute("ALTER TABLE user ADD COLUMN two_factor_enabled BOOLEAN DEFAULT 0") + + if 'two_factor_secret' not in user_columns: + print("Adding two_factor_secret column to user table...") + cursor.execute("ALTER TABLE user ADD COLUMN two_factor_secret VARCHAR(32)") + + # Check if the team table exists + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='team'") + if not cursor.fetchone(): + print("Creating team table...") + cursor.execute(""" + CREATE TABLE team ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name VARCHAR(100) UNIQUE NOT NULL, + description VARCHAR(255), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) # Check if the system_settings table exists cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='system_settings'") @@ -140,7 +172,9 @@ def migrate_database(): username='admin', email='admin@timetrack.local', is_admin=True, - is_verified=True # Admin is automatically verified + is_verified=True, # Admin is automatically verified + role=Role.ADMIN, + two_factor_enabled=False ) admin.set_password('admin') # Default password, should be changed db.session.add(admin) @@ -148,11 +182,15 @@ def migrate_database(): print("Created admin user with username 'admin' and password 'admin'") print("Please change the admin password after first login!") else: - # Make sure existing admin user is verified + # Make sure existing admin user is verified and has correct role if not hasattr(admin, 'is_verified') or not admin.is_verified: admin.is_verified = True - db.session.commit() - print("Marked existing admin user as verified") + if not hasattr(admin, 'role') or admin.role is None: + admin.role = Role.ADMIN + if not hasattr(admin, 'two_factor_enabled') or admin.two_factor_enabled is None: + admin.two_factor_enabled = False + db.session.commit() + print("Updated existing admin user with new fields") # Update existing time entries to associate with admin user orphan_entries = TimeEntry.query.filter_by(user_id=None).all() @@ -168,11 +206,29 @@ def migrate_database(): existing_users = User.query.filter_by(is_verified=None).all() for user in existing_users: user.is_verified = True + + # Update existing users with default role and 2FA settings + users_to_update = User.query.all() + updated_count = 0 + for user in users_to_update: + updated = False + if not hasattr(user, 'role') or user.role is None: + if user.is_admin: + user.role = Role.ADMIN + else: + user.role = Role.TEAM_MEMBER + updated = True + if not hasattr(user, 'two_factor_enabled') or user.two_factor_enabled is None: + user.two_factor_enabled = False + updated = True + if updated: + updated_count += 1 db.session.commit() print(f"Associated {len(orphan_entries)} existing time entries with admin user") print(f"Associated {len(orphan_configs)} existing work configs with admin user") print(f"Marked {len(existing_users)} existing users as verified") + print(f"Updated {updated_count} users with default role and 2FA settings") def init_system_settings(): """Initialize system settings with default values if they don't exist""" diff --git a/models.py b/models.py index 8ed7ba5..bdebfda 100644 --- a/models.py +++ b/models.py @@ -47,6 +47,10 @@ class User(db.Model): role = db.Column(db.Enum(Role), default=Role.TEAM_MEMBER) team_id = db.Column(db.Integer, db.ForeignKey('team.id'), nullable=True) + # Two-Factor Authentication fields + two_factor_enabled = db.Column(db.Boolean, default=False) + two_factor_secret = db.Column(db.String(32), nullable=True) # Base32 encoded secret + # Relationships time_entries = db.relationship('TimeEntry', backref='user', lazy=True) work_config = db.relationship('WorkConfig', backref='user', lazy=True, uselist=False) @@ -72,6 +76,34 @@ class User(db.Model): return True return False + def generate_2fa_secret(self): + """Generate a new 2FA secret""" + import pyotp + self.two_factor_secret = pyotp.random_base32() + return self.two_factor_secret + + def get_2fa_uri(self): + """Get the provisioning URI for QR code generation""" + if not self.two_factor_secret: + return None + import pyotp + totp = pyotp.TOTP(self.two_factor_secret) + return totp.provisioning_uri( + name=self.email, + issuer_name="TimeTrack" + ) + + def verify_2fa_token(self, token, allow_setup=False): + """Verify a 2FA token""" + if not self.two_factor_secret: + return False + # During setup, allow verification even if 2FA isn't enabled yet + if not allow_setup and not self.two_factor_enabled: + return False + import pyotp + totp = pyotp.TOTP(self.two_factor_secret) + return totp.verify(token, valid_window=1) # Allow 1 window tolerance + def __repr__(self): return f'' diff --git a/requirements.txt b/requirements.txt index ade5b39..2bf148d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,6 @@ itsdangerous==2.0.1 click==8.0.1 Flask-SQLAlchemy==2.5.1 SQLAlchemy==1.4.23 -python-dotenv==0.19.0 \ No newline at end of file +python-dotenv==0.19.0 +pyotp==2.6.0 +qrcode[pil]==7.3.1 \ No newline at end of file diff --git a/templates/profile.html b/templates/profile.html index b093036..82f2b5f 100644 --- a/templates/profile.html +++ b/templates/profile.html @@ -16,6 +16,13 @@

Username: {{ user.username }}

Account Type: {% if user.is_admin %}Administrator{% else %}User{% endif %}

Member Since: {{ user.created_at.strftime('%Y-%m-%d') }}

+

Two-Factor Authentication: + {% if user.two_factor_enabled %} + ✅ Enabled + {% else %} + ❌ Disabled + {% endif %} +

Update Profile

@@ -45,5 +52,107 @@ + +
+

Security Settings

+ +
+

Two-Factor Authentication

+ {% if user.two_factor_enabled %} +

Two-factor authentication is enabled for your account. This adds an extra layer of security by requiring a code from your authenticator app when logging in.

+ +
+
+ + +
+ +
+ {% else %} +

Two-factor authentication is not enabled for your account. We strongly recommend enabling it to protect your account.

+

With 2FA enabled, you'll need both your password and a code from your phone to log in.

+ + Enable Two-Factor Authentication + {% endif %} +
+
+ + {% endblock %} \ No newline at end of file diff --git a/templates/setup_2fa.html b/templates/setup_2fa.html new file mode 100644 index 0000000..e8b2791 --- /dev/null +++ b/templates/setup_2fa.html @@ -0,0 +1,266 @@ +{% extends "layout.html" %} + +{% block content %} +
+

Setup Two-Factor Authentication

+ +
+
+

Step 1: Install an Authenticator App

+

Download and install an authenticator app on your mobile device:

+
    +
  • Google Authenticator (iOS/Android)
    • +
  • Microsoft Authenticator (iOS/Android)
    • +
  • Authy (iOS/Android/Desktop)
    • +
  • 1Password (Premium feature)
    • +
+
+ +
+

Step 2: Scan QR Code or Enter Secret

+
+
+ 2FA QR Code +
+
+

Can't scan? Enter this code manually:

+
{{ secret }}
+

Account: {{ g.user.email }}
Issuer: TimeTrack

+
+
+
+ +
+

Step 3: Verify Setup

+

Enter the 6-digit code from your authenticator app to complete setup:

+ +
+
+ + + Enter the 6-digit code from your authenticator app +
+ +
+ + Cancel +
+
+
+
+ +
+

🔐 Security Notice

+

Important: Once enabled, you'll need your authenticator app to log in. Make sure to:

+
    +
  • Keep your authenticator app secure and backed up
    • +
  • Store the secret code in a safe place as a backup
    • +
  • Remember your password - you'll need both your password and 2FA code to log in
    • +
+
+
+ + + + +{% endblock %} \ No newline at end of file diff --git a/templates/verify_2fa.html b/templates/verify_2fa.html new file mode 100644 index 0000000..354558b --- /dev/null +++ b/templates/verify_2fa.html @@ -0,0 +1,182 @@ +{% extends "layout.html" %} + +{% block content %} +
+
+

Two-Factor Authentication

+

Please enter the 6-digit code from your authenticator app to complete login.

+ +
+
+ + + Enter the 6-digit code from your authenticator app +
+ +
+ +
+
+ +
+

Having trouble? Make sure your device's time is synchronized and try a new code.

+

← Back to Login

+
+
+
+ + + + +{% endblock %} \ No newline at end of file