171 lines
6.9 KiB
Python
171 lines
6.9 KiB
Python
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'))
|