import io import base64 import logging import pyotp import qrcode from flask import (Blueprint, render_template, redirect, url_for, flash, request, session) from flask_login import login_user, logout_user, login_required, current_user from flask_wtf import FlaskForm from wtforms import StringField, PasswordField, BooleanField, SubmitField from wtforms.validators import DataRequired, Length from app.models.user import User from app.extensions import db, limiter from datetime import datetime auth_bp = Blueprint('auth', __name__, url_prefix='/auth') log = logging.getLogger(__name__) class LoginForm(FlaskForm): username = StringField('Username', validators=[DataRequired(), Length(1, 64)]) password = PasswordField('Password', validators=[DataRequired()]) remember_me = BooleanField('Remember me') submit = SubmitField('Sign In') # ── Login / logout ──────────────────────────────────────────────────────────── @auth_bp.route('/login', methods=['GET', 'POST']) @limiter.limit('10 per minute; 30 per hour') def login(): if current_user.is_authenticated: return redirect(url_for('dashboard.index')) form = LoginForm() if form.validate_on_submit(): user = User.query.filter_by(username=form.username.data.strip()).first() if user and user.check_password(form.password.data): if user.totp_enabled: # Park the user id in the session and redirect to TOTP verification session['_totp_pending_id'] = user.id session['_totp_remember'] = bool(form.remember_me.data) session['_totp_next'] = request.args.get('next', '') log.info('[auth] TOTP required for user %s', user.username) return redirect(url_for('auth.totp_verify')) login_user(user, remember=form.remember_me.data) user.last_login = datetime.utcnow() db.session.commit() log.info('[auth] user %s logged in', user.username) next_page = request.args.get('next', '') if not next_page.startswith('/'): next_page = url_for('dashboard.index') return redirect(next_page) log.warning('[auth] failed login attempt for username=%r ip=%s', form.username.data, request.remote_addr) flash('Invalid username or password.', 'danger') return render_template('auth/login.html', form=form) @auth_bp.route('/logout') @login_required def logout(): logout_user() flash('You have been logged out.', 'info') return redirect(url_for('auth.login')) # ── TOTP: second-factor verification ───────────────────────────────────────── @auth_bp.route('/totp/verify', methods=['GET', 'POST']) @limiter.limit('10 per minute') def totp_verify(): pending_id = session.get('_totp_pending_id') if not pending_id: return redirect(url_for('auth.login')) user = db.session.get(User, pending_id) if not user or not user.totp_enabled: session.pop('_totp_pending_id', None) return redirect(url_for('auth.login')) error = None if request.method == 'POST': code = request.form.get('code', '').strip().replace(' ', '') totp = pyotp.TOTP(user.totp_secret) if totp.verify(code, valid_window=1): session.pop('_totp_pending_id', None) remember = session.pop('_totp_remember', False) next_url = session.pop('_totp_next', '') or url_for('dashboard.index') if not next_url.startswith('/'): next_url = url_for('dashboard.index') login_user(user, remember=remember) user.last_login = datetime.utcnow() db.session.commit() log.info('[auth] TOTP verified for user %s', user.username) return redirect(next_url) log.warning('[auth] invalid TOTP code for user %s ip=%s', user.username, request.remote_addr) error = 'Invalid code — please try again.' return render_template('auth/totp_verify.html', error=error) # ── TOTP: setup (enable) ────────────────────────────────────────────────────── @auth_bp.route('/totp/setup', methods=['GET', 'POST']) @login_required def totp_setup(): user = current_user if request.method == 'GET': # Generate a new secret on every GET so the user always sees a fresh QR code secret = pyotp.random_base32() session['_totp_setup_secret'] = secret totp = pyotp.TOTP(secret) uri = totp.provisioning_uri( name=user.username, issuer_name='PFM — Personal Finance', ) img = qrcode.make(uri) buf = io.BytesIO() img.save(buf, format='PNG') buf.seek(0) qr_data = 'data:image/png;base64,' + base64.b64encode(buf.read()).decode() return render_template('auth/totp_setup.html', qr_data=qr_data, secret=secret, uri=uri) # POST — user confirmed with their first code secret = session.get('_totp_setup_secret') code = request.form.get('code', '').strip().replace(' ', '') if not secret: flash('Setup session expired. Please try again.', 'warning') return redirect(url_for('auth.totp_setup')) totp = pyotp.TOTP(secret) if totp.verify(code, valid_window=1): user.totp_secret = secret user.totp_enabled = True db.session.commit() session.pop('_totp_setup_secret', None) log.info('[auth] TOTP enabled for user %s', user.username) flash('Two-factor authentication enabled successfully.', 'success') return redirect(url_for('settings.index')) flash('Invalid code — the QR code was not scanned correctly. Please try again.', 'danger') return redirect(url_for('auth.totp_setup')) # ── TOTP: disable ───────────────────────────────────────────────────────────── @auth_bp.route('/totp/disable', methods=['POST']) @login_required def totp_disable(): user = current_user # Require password confirmation before disabling password = request.form.get('password', '') if not user.check_password(password): flash('Incorrect password — 2FA not disabled.', 'danger') return redirect(url_for('settings.index')) user.totp_enabled = False user.totp_secret = None db.session.commit() log.info('[auth] TOTP disabled for user %s', user.username) flash('Two-factor authentication disabled.', 'info') return redirect(url_for('settings.index'))