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 app.utils.audit import audit 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() audit('login_success', f'user={user.username}') 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) audit('login_failed', f'username={form.username.data!r}') 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() audit('login_success_2fa', f'user={user.username}') 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) audit('totp_enabled', f'user={user.username}') 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() audit('totp_disabled', f'user={user.username}') log.info('[auth] TOTP disabled for user %s', user.username) flash('Two-factor authentication disabled.', 'info') return redirect(url_for('settings.index'))