115 lines
4 KiB
Python
Executable file
115 lines
4 KiB
Python
Executable file
from flask import Blueprint, render_template, request, redirect, url_for, flash, session
|
|
from flask_login import login_user, logout_user, login_required, current_user
|
|
from app import db
|
|
from app.models.user import User
|
|
|
|
bp = Blueprint('auth', __name__, url_prefix='/auth')
|
|
|
|
@bp.route('/login', methods=['GET', 'POST'])
|
|
def login():
|
|
if current_user.is_authenticated:
|
|
return redirect(url_for('main.dashboard'))
|
|
|
|
if request.method == 'POST':
|
|
username = request.form.get('username')
|
|
password = request.form.get('password')
|
|
totp_token = request.form.get('totp_token')
|
|
|
|
user = User.query.filter_by(username=username).first()
|
|
|
|
if user and user.check_password(password):
|
|
# Check if 2FA is enabled
|
|
if user.is_2fa_enabled:
|
|
if not totp_token:
|
|
# Store user ID in session for 2FA verification
|
|
session['2fa_user_id'] = user.id
|
|
return render_template('auth/verify_2fa.html')
|
|
else:
|
|
# Verify 2FA token
|
|
if user.verify_totp(totp_token):
|
|
login_user(user)
|
|
session.pop('2fa_user_id', None)
|
|
flash('Login successful!', 'success')
|
|
return redirect(url_for('main.dashboard'))
|
|
else:
|
|
flash('Invalid 2FA code', 'error')
|
|
return render_template('auth/verify_2fa.html')
|
|
else:
|
|
# No 2FA, login directly
|
|
login_user(user)
|
|
flash('Login successful!', 'success')
|
|
return redirect(url_for('main.dashboard'))
|
|
|
|
flash('Invalid username or password', 'error')
|
|
|
|
return render_template('auth/login.html')
|
|
|
|
@bp.route('/verify-2fa', methods=['POST'])
|
|
def verify_2fa():
|
|
user_id = session.get('2fa_user_id')
|
|
if not user_id:
|
|
flash('Session expired. Please login again.', 'error')
|
|
return redirect(url_for('auth.login'))
|
|
|
|
user = User.query.get(user_id)
|
|
if not user:
|
|
flash('User not found', 'error')
|
|
return redirect(url_for('auth.login'))
|
|
|
|
token = request.form.get('token')
|
|
|
|
if user.verify_totp(token):
|
|
login_user(user)
|
|
session.pop('2fa_user_id', None)
|
|
flash('Login successful!', 'success')
|
|
return redirect(url_for('main.dashboard'))
|
|
else:
|
|
flash('Invalid 2FA code. Please try again.', 'error')
|
|
return render_template('auth/verify_2fa.html')
|
|
|
|
@bp.route('/register', methods=['GET', 'POST'])
|
|
def register():
|
|
if current_user.is_authenticated:
|
|
return redirect(url_for('main.dashboard'))
|
|
|
|
if request.method == 'POST':
|
|
username = request.form.get('username')
|
|
email = request.form.get('email')
|
|
password = request.form.get('password')
|
|
confirm_password = request.form.get('confirm_password')
|
|
|
|
if password != confirm_password:
|
|
flash('Passwords do not match', 'error')
|
|
return redirect(url_for('auth.register'))
|
|
|
|
if User.query.filter_by(username=username).first():
|
|
flash('Username already exists', 'error')
|
|
return redirect(url_for('auth.register'))
|
|
|
|
if User.query.filter_by(email=email).first():
|
|
flash('Email already exists', 'error')
|
|
return redirect(url_for('auth.register'))
|
|
|
|
is_first_user = User.query.count() == 0
|
|
|
|
user = User(
|
|
username=username,
|
|
email=email,
|
|
is_admin=is_first_user
|
|
)
|
|
user.set_password(password)
|
|
|
|
db.session.add(user)
|
|
db.session.commit()
|
|
|
|
flash('Registration successful! Please login.', 'success')
|
|
return redirect(url_for('auth.login'))
|
|
|
|
return render_template('auth/register.html')
|
|
|
|
@bp.route('/logout')
|
|
@login_required
|
|
def logout():
|
|
logout_user()
|
|
flash('Logged out successfully', 'success')
|
|
return redirect(url_for('auth.login'))
|