Infosec/Alert Hub System/app/routes_auth.py

"""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"))