#!/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 = 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 | tuple[str, int]: """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") token = flask.request.args.get("token") if not token: return "Invalid token", 400 json_data = UniAuth.auth.verify_secure_token( token, salt="secure-redirect", max_age=600 ) assert isinstance(json_data, str) token_payload = json.loads(json_data) if not token_payload: return "Invalid token", 400 callback_url = token_payload["callback_url"] auth_token = UniAuth.auth.generate_auth_token(user["username"]) redirect_to_callback = f"{callback_url}?auth_token={auth_token}&next={token}" # 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( "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")