UniAuth/main.py

#!/usr/bin/python3
"""Single sign-on application."""

import functools
import json
import typing
from datetime import datetime, timedelta
from typing import Any, Callable

import flask
import werkzeug

import UniAuth.auth

app = flask.Flask(__name__)
app.debug = False
app.config.from_object("config.default")


def login_required(f: Callable[..., Any]) -> Callable[..., Any]:
    """Route requires login decorator."""

    @functools.wraps(f)
    def decorated_function(*args, **kwargs) -> werkzeug.Response:  # type: ignore
        token = flask.request.cookies.get("uniauth_token")
        if not token or UniAuth.auth.verify_auth_token(token) is None:
            # Save the original URL in the session and redirect to login
            flask.session["next"] = flask.request.url
            return flask.redirect(flask.url_for("login_page"))
        return typing.cast(werkzeug.Response, f(*args, **kwargs))

    return decorated_function


@app.route("/")
@login_required
def root_page() -> str:
    """Root page."""
    return flask.render_template("auth_good.html")


def check_user_auth() -> dict[str, Any] | None:
    """Load username and password and check if valid."""
    # Extract username and password from POST request
    username = flask.request.form.get("username", type=str)
    password = flask.request.form.get("password", type=str)

    # Retrieve users from app config
    users = app.config["USERS"]

    # Check if the username and password match any user in the list
    return next(
        (u for u in users if u["username"] == username and u["password"] == password),
        None,
    )


def read_callback_url_from_token() -> str | None:
    """Parse token and extract callback URL."""
    token = flask.request.args.get("token")
    if not token:
        return None

    json_data = UniAuth.auth.verify_secure_token(
        token, salt="secure-redirect", max_age=600
    )
    if not json_data:
        return None
    assert isinstance(json_data, str)
    token_payload = json.loads(json_data)

    callback_url = token_payload["callback_url"]
    assert isinstance(callback_url, str)

    return callback_url


@app.route("/login", methods=["GET", "POST"])
def login_page() -> str | werkzeug.Response | tuple[str, int]:
    """Login page."""
    app.logger.info("Login page.")
    if flask.request.method == "GET":
        uniauth_token = flask.request.cookies.get("uniauth_token")
        if (
            not uniauth_token
            or not UniAuth.auth.verify_auth_token(uniauth_token)
            or not (callback_url := read_callback_url_from_token())
        ):
            return flask.render_template("login.html")
        token = flask.request.args["token"]
        redirect_to_callback = f"{callback_url}?auth_token={uniauth_token}&next={token}"
        app.logger.info(f"Redirecting to: {redirect_to_callback}")

        return flask.redirect(redirect_to_callback)

    if not (user := check_user_auth()):
        # Login failed: Show an error message on the login page
        app.logger.info("User auth failed")
        return flask.render_template("login.html", error="Invalid credentials")

    callback_url = read_callback_url_from_token()
    if not callback_url:
        return "Invalid token", 400

    auth_token = UniAuth.auth.generate_auth_token(user["username"])
    token = flask.request.args["token"]

    redirect_to_callback = f"{callback_url}?auth_token={auth_token}&next={token}"

    app.logger.info(f"Redirecting to: {redirect_to_callback}")

    # flask.flash("Welcome back! You have successfully logged in.")

    expire_date = datetime.now() + timedelta(days=180)

    response = flask.redirect(redirect_to_callback)
    response.set_cookie(
        "uniauth_token",
        token,
        expires=expire_date,
        httponly=True,
        secure=True,
        samesite="Lax",
    )

    return response


@app.route("/logout")
def logout() -> werkzeug.Response:
    """Handle user logout by clearing the authentication cookie."""
    after_login = flask.request.args.get("next")
    response = flask.redirect(flask.url_for("login_page", next=after_login))
    response.set_cookie("auth_token", "", expires=0)
    flask.flash("You have been successfully logged out.", "info")
    return response


if __name__ == "__main__":
    app.run(host="0.0.0.0")