#!/usr/bin/python3 """Single sign-on application.""" import functools 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 = True 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("auth_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, ) @app.route("/login", methods=["GET", "POST"]) def login_page() -> str | werkzeug.Response: """Login page.""" if flask.request.method == "GET": return flask.render_template("login.html") if not (user := check_user_auth()): # Login failed: Show an error message on the login page return flask.render_template("login.html", error="Invalid credentials") redirect_to = ( flask.request.args.get("next") or flask.session.get("next") or flask.url_for("dashboard") ) expire_date = datetime.now() + timedelta(days=180) flask.flash("Welcome back! You have successfully logged in.") token = UniAuth.auth.generate_auth_token(user["username"]) response = flask.redirect(redirect_to) response.set_cookie( "auth_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")