v_1.5
This commit is contained in:
jpmvaz
@@ -0,0 +1,308 @@
|
||||
"""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/<int:user_id>", 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"))
|
||||
Reference in New Issue
Block a user