#!/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 from itsdangerous.url_safe import URLSafeTimedSerializer app = flask.Flask(__name__) app.debug = True app.config.from_object("config.default") serializer = URLSafeTimedSerializer(app.config["SECRET_KEY"]) max_age = 60 * 60 * 24 * 90 def generate_auth_token(username: str) -> str: """Generate a secure authentication token.""" token = typing.cast(str, serializer.dumps(username, salt="auth")) assert isinstance(token, str) return token def verify_auth_token(token: str) -> str | None: """Verify the authentication token.""" try: username = serializer.loads(token, salt="auth", max_age=max_age) except Exception: return None assert isinstance(username, str) return username 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 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.") response = flask.redirect(redirect_to) response.set_cookie( "auth_token", generate_auth_token(user["username"]), 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.""" response = flask.redirect(flask.url_for("login_page")) 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")