"""Authentication routes: first-run setup, login, MFA, logout, password, avatar.""" import io import base64 from datetime import timedelta import pyotp import qrcode from PIL import Image, ImageOps from flask import ( Blueprint, render_template, redirect, url_for, flash, request, session, current_app, abort, Response, ) from flask_login import login_user, logout_user, login_required, current_user from .models import db, User, AuditLog, utcnow from .forms import LoginForm, MfaForm, MfaSetupForm, ChangePasswordForm, SetupAdminForm, AvatarForm from .auth_utils import get_remote_ip, login_required_full auth_bp = Blueprint("auth", __name__) AVATAR_SIZE = 256 def _normalise_avatar(file_storage) -> bytes: """Open an uploaded image, square-crop + resize, return PNG bytes. Raises ValueError if the file is not a usable image. """ try: img = Image.open(file_storage.stream) img = ImageOps.exif_transpose(img) img = img.convert("RGBA") if img.mode in ("RGBA", "LA", "P") else img.convert("RGB") # Center-crop to a square, then resize. img = ImageOps.fit(img, (AVATAR_SIZE, AVATAR_SIZE), method=Image.LANCZOS) if img.mode == "RGBA": background = Image.new("RGBA", img.size, (255, 255, 255, 0)) img = Image.alpha_composite(background, img) buf = io.BytesIO() img.save(buf, format="PNG", optimize=True) return buf.getvalue() except Exception as exc: # noqa: BLE001 - surface a clean error to the user raise ValueError(f"Could not process image: {exc}") # --------------------------------------------------------------------------- # # First-run setup — create the initial administrator # --------------------------------------------------------------------------- # @auth_bp.route("/setup", methods=["GET", "POST"]) def setup(): # Only available while there are zero users. Once an admin exists, hide it. if User.query.count() > 0: return redirect(url_for("auth.login")) form = SetupAdminForm() if form.validate_on_submit(): admin = User( username=form.username.data.strip(), email=form.email.data.strip().lower(), role="admin", is_active=True, ) admin.set_password(form.password.data) db.session.add(admin) db.session.commit() current_app._alerthub_has_users = True AuditLog.record("setup.admin_created", user=admin, details=f"username={admin.username}", ip=get_remote_ip()) flash("Administrator account created. Please sign in.", "success") return redirect(url_for("auth.login")) return render_template("setup.html", form=form) # --------------------------------------------------------------------------- # # Login # --------------------------------------------------------------------------- # @auth_bp.route("/login", methods=["GET", "POST"]) def login(): if current_user.is_authenticated and not session.get("mfa_pending"): return redirect(url_for("main.index")) form = LoginForm() if form.validate_on_submit(): user = User.query.filter_by(username=form.username.data.strip()).first() ip = get_remote_ip() if user is None or not user.check_password(form.password.data): # Increment failed login if user exists if user is not None: user.failed_login_count = (user.failed_login_count or 0) + 1 if user.failed_login_count >= 5: user.locked_until = utcnow() + timedelta(minutes=15) user.failed_login_count = 0 AuditLog.record("login.lockout", user=user, details=f"5 failed attempts", ip=ip) db.session.commit() AuditLog.record("login.fail", details=f"username={form.username.data}", ip=ip) flash("Invalid credentials.", "error") return render_template("login.html", form=form), 401 if not user.is_active: flash("Account is disabled.", "error") AuditLog.record("login.disabled", user=user, ip=ip) return render_template("login.html", form=form), 403 if user.is_locked: flash("Account is temporarily locked due to failed login attempts. Try again later.", "error") AuditLog.record("login.locked", user=user, ip=ip) return render_template("login.html", form=form), 423 # Password OK — reset counters user.failed_login_count = 0 user.locked_until = None if user.mfa_enabled: # Stage 1 success — but require MFA before granting full session session["mfa_pending"] = True session["pending_user_id"] = user.id db.session.commit() AuditLog.record("login.password_ok", user=user, ip=ip) return redirect(url_for("auth.mfa_challenge")) # No MFA — full login login_user(user, remember=False) session.permanent = True session.pop("mfa_pending", None) session.pop("pending_user_id", None) user.last_login_at = utcnow() db.session.commit() AuditLog.record("login.success", user=user, ip=ip) next_url = request.args.get("next") or url_for("main.index") return redirect(next_url) return render_template("login.html", form=form) # --------------------------------------------------------------------------- # # MFA challenge (after password OK) # --------------------------------------------------------------------------- # @auth_bp.route("/mfa", methods=["GET", "POST"]) def mfa_challenge(): if not session.get("mfa_pending"): return redirect(url_for("auth.login")) user_id = session.get("pending_user_id") user = db.session.get(User, user_id) if user_id else None if user is None or not user.mfa_enabled or not user.mfa_secret: session.pop("mfa_pending", None) session.pop("pending_user_id", None) return redirect(url_for("auth.login")) form = MfaForm() if form.validate_on_submit(): totp = pyotp.TOTP(user.mfa_secret) if totp.verify(form.code.data, valid_window=1): login_user(user, remember=False) session.permanent = True session.pop("mfa_pending", None) session.pop("pending_user_id", None) user.last_login_at = utcnow() db.session.commit() AuditLog.record("login.mfa_ok", user=user, ip=get_remote_ip()) return redirect(url_for("main.index")) flash("Invalid authentication code.", "error") AuditLog.record("login.mfa_fail", user=user, ip=get_remote_ip()) return render_template("mfa.html", form=form) # --------------------------------------------------------------------------- # # Logout # --------------------------------------------------------------------------- # @auth_bp.route("/logout", methods=["POST"]) @login_required def logout(): AuditLog.record("logout", user=current_user, ip=get_remote_ip()) logout_user() session.clear() flash("Signed out.", "success") return redirect(url_for("auth.login")) # --------------------------------------------------------------------------- # # Account: change password / set up MFA # --------------------------------------------------------------------------- # @auth_bp.route("/account", methods=["GET"]) @login_required_full def account(): return render_template( "account.html", password_form=ChangePasswordForm(), mfa_setup_form=MfaSetupForm(), avatar_form=AvatarForm(), ) # --------------------------------------------------------------------------- # # Avatar: upload / remove / serve # --------------------------------------------------------------------------- # @auth_bp.route("/account/avatar", methods=["POST"]) @login_required_full def avatar_upload(): form = AvatarForm() if form.validate_on_submit(): try: png_bytes = _normalise_avatar(form.avatar.data) except ValueError as e: flash(str(e), "error") return redirect(url_for("auth.account")) current_user.avatar_data = png_bytes current_user.avatar_mime = "image/png" db.session.commit() AuditLog.record("avatar.update", user=current_user, ip=get_remote_ip()) flash("Avatar updated.", "success") else: for errs in form.errors.values(): for e in errs: flash(e, "error") return redirect(url_for("auth.account")) @auth_bp.route("/account/avatar/remove", methods=["POST"]) @login_required_full def avatar_remove(): current_user.avatar_data = None current_user.avatar_mime = None db.session.commit() AuditLog.record("avatar.remove", user=current_user, ip=get_remote_ip()) flash("Avatar removed.", "info") return redirect(url_for("auth.account")) @auth_bp.route("/avatar/", methods=["GET"]) @login_required_full def avatar(user_id: int): user = db.session.get(User, user_id) if user is None or not user.has_avatar: abort(404) resp = Response(user.avatar_data, mimetype=user.avatar_mime or "image/png") resp.headers["Cache-Control"] = "private, max-age=60" return resp @auth_bp.route("/account/password", methods=["POST"]) @login_required_full def change_password(): form = ChangePasswordForm() if form.validate_on_submit(): if not current_user.check_password(form.current_password.data): flash("Current password is incorrect.", "error") else: current_user.set_password(form.new_password.data) db.session.commit() AuditLog.record("password.change", user=current_user, ip=get_remote_ip()) flash("Password updated.", "success") return redirect(url_for("auth.account")) return render_template( "account.html", password_form=form, mfa_setup_form=MfaSetupForm(), ) @auth_bp.route("/account/mfa/setup", methods=["GET", "POST"]) @login_required_full def mfa_setup(): """Generate a new TOTP secret + QR code; user confirms with a 6-digit code.""" if current_user.mfa_enabled: flash("MFA is already enabled.", "info") return redirect(url_for("auth.account")) # Keep candidate secret in session until verified if "mfa_setup_secret" not in session: session["mfa_setup_secret"] = pyotp.random_base32() secret = session["mfa_setup_secret"] issuer = current_app.config.get("MFA_ISSUER_NAME", "AlertHub") uri = pyotp.TOTP(secret).provisioning_uri(name=current_user.email, issuer_name=issuer) # Generate QR code as base64 PNG img = qrcode.make(uri) buf = io.BytesIO() img.save(buf, format="PNG") qr_b64 = base64.b64encode(buf.getvalue()).decode("ascii") form = MfaSetupForm() if form.validate_on_submit(): if pyotp.TOTP(secret).verify(form.code.data, valid_window=1): current_user.mfa_secret = secret current_user.mfa_enabled = True db.session.commit() session.pop("mfa_setup_secret", None) AuditLog.record("mfa.enabled", user=current_user, ip=get_remote_ip()) flash("MFA enabled successfully. You will be prompted for a code on next sign-in.", "success") return redirect(url_for("auth.account")) flash("Verification code did not match. Try again.", "error") return render_template("mfa_setup.html", form=form, qr_b64=qr_b64, secret=secret) @auth_bp.route("/account/mfa/disable", methods=["POST"]) @login_required_full def mfa_disable(): current_user.mfa_enabled = False current_user.mfa_secret = None db.session.commit() AuditLog.record("mfa.disabled", user=current_user, ip=get_remote_ip()) flash("MFA disabled.", "info") return redirect(url_for("auth.account"))