from flask import Blueprint, render_template, redirect, url_for, flash, request, session, send_file, make_response from flask_login import login_user, logout_user, login_required, current_user from app import db, bcrypt from app.models import User import pyotp import qrcode import io import base64 import secrets import json from datetime import datetime bp = Blueprint('auth', __name__, url_prefix='/auth') def generate_backup_codes(count=10): """Generate backup codes for 2FA""" codes = [] for _ in range(count): # Generate 8-character alphanumeric code code = ''.join(secrets.choice('ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789') for _ in range(8)) # Format as XXXX-XXXX for readability formatted_code = f"{code[:4]}-{code[4:]}" codes.append(formatted_code) return codes def hash_backup_codes(codes): """Hash backup codes for secure storage""" return [bcrypt.generate_password_hash(code).decode('utf-8') for code in codes] def verify_backup_code(user, code): """Verify a backup code and mark it as used""" if not user.backup_codes: return False stored_codes = json.loads(user.backup_codes) for i, hashed_code in enumerate(stored_codes): if bcrypt.check_password_hash(hashed_code, code): # Remove used code stored_codes.pop(i) user.backup_codes = json.dumps(stored_codes) db.session.commit() return True return False @bp.route('/login', methods=['GET', 'POST']) def login(): if current_user.is_authenticated: return redirect(url_for('main.dashboard')) if request.method == 'POST': data = request.get_json() if request.is_json else request.form username = data.get('username') password = data.get('password') two_factor_code = data.get('two_factor_code') remember = data.get('remember', False) # Accept both username and email user = User.query.filter((User.username == username) | (User.email == username)).first() if user and bcrypt.check_password_hash(user.password_hash, password): # Check 2FA if enabled if user.two_factor_enabled: if not two_factor_code: if request.is_json: return {'success': False, 'requires_2fa': True}, 200 session['pending_user_id'] = user.id return render_template('auth/two_factor.html') # Try TOTP code first totp = pyotp.TOTP(user.totp_secret) is_valid = totp.verify(two_factor_code) # If TOTP fails, try backup code (format: XXXX-XXXX or XXXXXXXX) if not is_valid: is_valid = verify_backup_code(user, two_factor_code) if not is_valid: if request.is_json: return {'success': False, 'message': 'Invalid 2FA code'}, 401 flash('Invalid 2FA code', 'error') return render_template('auth/login.html') login_user(user, remember=remember) session.permanent = remember if request.is_json: return {'success': True, 'redirect': url_for('main.dashboard')} next_page = request.args.get('next') return redirect(next_page if next_page else url_for('main.dashboard')) if request.is_json: return {'success': False, 'message': 'Invalid username or password'}, 401 flash('Invalid username or password', 'error') return render_template('auth/login.html') @bp.route('/register', methods=['GET', 'POST']) def register(): if current_user.is_authenticated: return redirect(url_for('main.dashboard')) if request.method == 'POST': data = request.get_json() if request.is_json else request.form username = data.get('username') email = data.get('email') password = data.get('password') language = data.get('language', 'en') currency = data.get('currency', 'USD') # Check if user exists if User.query.filter_by(email=email).first(): if request.is_json: return {'success': False, 'message': 'Email already registered'}, 400 flash('Email already registered', 'error') return render_template('auth/register.html') if User.query.filter_by(username=username).first(): if request.is_json: return {'success': False, 'message': 'Username already taken'}, 400 flash('Username already taken', 'error') return render_template('auth/register.html') # Check if this is the first user (make them admin) is_first_user = User.query.count() == 0 # Create user password_hash = bcrypt.generate_password_hash(password).decode('utf-8') user = User( username=username, email=email, password_hash=password_hash, is_admin=is_first_user, language=language, currency=currency ) db.session.add(user) db.session.commit() # Create default categories from app.utils import create_default_categories create_default_categories(user.id) login_user(user) if request.is_json: return {'success': True, 'redirect': url_for('main.dashboard')} flash('Registration successful!', 'success') return redirect(url_for('main.dashboard')) return render_template('auth/register.html') @bp.route('/logout') @login_required def logout(): logout_user() return redirect(url_for('auth.login')) @bp.route('/setup-2fa', methods=['GET', 'POST']) @login_required def setup_2fa(): if request.method == 'POST': data = request.get_json() if request.is_json else request.form code = data.get('code') if not current_user.totp_secret: secret = pyotp.random_base32() current_user.totp_secret = secret totp = pyotp.TOTP(current_user.totp_secret) if totp.verify(code): # Generate backup codes backup_codes_plain = generate_backup_codes(10) backup_codes_hashed = hash_backup_codes(backup_codes_plain) current_user.two_factor_enabled = True current_user.backup_codes = json.dumps(backup_codes_hashed) db.session.commit() # Store plain backup codes in session for display session['backup_codes'] = backup_codes_plain if request.is_json: return {'success': True, 'message': '2FA enabled successfully', 'backup_codes': backup_codes_plain} flash('2FA enabled successfully', 'success') return redirect(url_for('auth.show_backup_codes')) if request.is_json: return {'success': False, 'message': 'Invalid code'}, 400 flash('Invalid code', 'error') # Generate QR code if not current_user.totp_secret: current_user.totp_secret = pyotp.random_base32() db.session.commit() totp = pyotp.TOTP(current_user.totp_secret) provisioning_uri = totp.provisioning_uri( name=current_user.email, issuer_name='FINA' ) # Generate QR code image qr = qrcode.QRCode(version=1, box_size=10, border=5) qr.add_data(provisioning_uri) qr.make(fit=True) img = qr.make_image(fill_color="black", back_color="white") buf = io.BytesIO() img.save(buf, format='PNG') buf.seek(0) qr_code_base64 = base64.b64encode(buf.getvalue()).decode() return render_template('auth/setup_2fa.html', qr_code=qr_code_base64, secret=current_user.totp_secret) @bp.route('/backup-codes', methods=['GET']) @login_required def show_backup_codes(): """Display backup codes after 2FA setup""" backup_codes = session.get('backup_codes', []) if not backup_codes: flash('No backup codes available', 'error') return redirect(url_for('main.settings')) return render_template('auth/backup_codes.html', backup_codes=backup_codes, username=current_user.username) @bp.route('/backup-codes/download', methods=['GET']) @login_required def download_backup_codes_pdf(): """Download backup codes as PDF""" backup_codes = session.get('backup_codes', []) if not backup_codes: flash('No backup codes available', 'error') return redirect(url_for('main.settings')) try: from reportlab.lib.pagesizes import letter from reportlab.lib.units import inch from reportlab.pdfgen import canvas from reportlab.lib import colors # Create PDF in memory buffer = io.BytesIO() c = canvas.Canvas(buffer, pagesize=letter) width, height = letter # Title c.setFont("Helvetica-Bold", 24) c.drawCentredString(width/2, height - 1*inch, "FINA") c.setFont("Helvetica-Bold", 18) c.drawCentredString(width/2, height - 1.5*inch, "Two-Factor Authentication") c.drawCentredString(width/2, height - 1.9*inch, "Backup Codes") # User info c.setFont("Helvetica", 12) c.drawString(1*inch, height - 2.5*inch, f"User: {current_user.username}") c.drawString(1*inch, height - 2.8*inch, f"Email: {current_user.email}") c.drawString(1*inch, height - 3.1*inch, f"Generated: {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}") # Warning message c.setFillColorRGB(0.8, 0.2, 0.2) c.setFont("Helvetica-Bold", 11) c.drawString(1*inch, height - 3.7*inch, "IMPORTANT: Store these codes in a secure location!") c.setFillColorRGB(0, 0, 0) c.setFont("Helvetica", 10) c.drawString(1*inch, height - 4.0*inch, "Each code can only be used once. Use them if you lose access to your authenticator app.") # Backup codes in two columns c.setFont("Courier-Bold", 14) y_position = height - 4.8*inch x_left = 1.5*inch x_right = 4.5*inch for i, code in enumerate(backup_codes): if i % 2 == 0: c.drawString(x_left, y_position, f"{i+1:2d}. {code}") else: c.drawString(x_right, y_position, f"{i+1:2d}. {code}") y_position -= 0.4*inch # Footer c.setFont("Helvetica", 8) c.setFillColorRGB(0.5, 0.5, 0.5) c.drawCentredString(width/2, 0.5*inch, "Keep this document secure and do not share these codes with anyone.") c.save() buffer.seek(0) # Clear backup codes from session after download session.pop('backup_codes', None) # Create response with PDF response = make_response(buffer.getvalue()) response.headers['Content-Type'] = 'application/pdf' response.headers['Content-Disposition'] = f'attachment; filename=FINA_BackupCodes_{current_user.username}_{datetime.utcnow().strftime("%Y%m%d")}.pdf' return response except ImportError: # If reportlab is not installed, return codes as text file text_content = f"FINA - Two-Factor Authentication Backup Codes\n\n" text_content += f"User: {current_user.username}\n" text_content += f"Email: {current_user.email}\n" text_content += f"Generated: {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')}\n\n" text_content += "IMPORTANT: Store these codes in a secure location!\n" text_content += "Each code can only be used once.\n\n" text_content += "Backup Codes:\n" text_content += "-" * 40 + "\n" for i, code in enumerate(backup_codes, 1): text_content += f"{i:2d}. {code}\n" text_content += "-" * 40 + "\n" text_content += "\nKeep this document secure and do not share these codes with anyone." # Clear backup codes from session session.pop('backup_codes', None) response = make_response(text_content) response.headers['Content-Type'] = 'text/plain' response.headers['Content-Disposition'] = f'attachment; filename=FINA_BackupCodes_{current_user.username}_{datetime.utcnow().strftime("%Y%m%d")}.txt' return response @bp.route('/disable-2fa', methods=['POST']) @login_required def disable_2fa(): current_user.two_factor_enabled = False current_user.backup_codes = None db.session.commit() if request.is_json: return {'success': True, 'message': '2FA disabled'} flash('2FA disabled', 'success') return redirect(url_for('main.settings'))