#!/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")