Personal-Finance-Management/app/routes/auth.py

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; 30 per hour')
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)
        attempts = session.get('_totp_attempts', 0) + 1
        session['_totp_attempts'] = attempts
        if attempts >= 5:
            session.pop('_totp_pending_id', None)
            session.pop('_totp_attempts', None)
            audit('login_failed_2fa', f'user={user.username} — max attempts reached')
            flash('Too many failed attempts. Please log in again.', 'danger')
            return redirect(url_for('auth.login'))
        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
@limiter.limit('10 per minute')
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'))