diff --git a/UniAuth/auth.py b/UniAuth/auth.py index 62c3a34..20f8446 100644 --- a/UniAuth/auth.py +++ b/UniAuth/auth.py @@ -1,6 +1,8 @@ """Authentication via UniAuth.""" +import json import typing +from datetime import datetime, timedelta import flask import werkzeug @@ -9,6 +11,31 @@ from itsdangerous.url_safe import URLSafeTimedSerializer max_age = 60 * 60 * 24 * 90 +def generate_secure_token(data: str, salt: str) -> str: + """Generate a secure token for the given data.""" + serializer = URLSafeTimedSerializer(flask.current_app.config["SECRET_KEY"]) + token = typing.cast(str, serializer.dumps(data, salt=salt)) + assert isinstance(token, str) + return token + + +def generate_auth_token(username: str) -> str: + """Generate a secure authentication token.""" + return generate_secure_token(username, "auth") + + +def verify_secure_token(token: str, salt: str, max_age: int) -> str | None: + """Verify the secure token and return the data if valid.""" + serializer = URLSafeTimedSerializer(flask.current_app.config["SECRET_KEY"]) + try: + data = serializer.loads(token, salt=salt, max_age=max_age) + except Exception: + return None + + assert isinstance(data, str) + return data + + def is_authentication() -> bool: """User is authenticated.""" return not flask.current_app.config.get("REQUIRE_AUTH") or bool( @@ -18,22 +45,7 @@ def is_authentication() -> bool: def verify_auth_token(token: str) -> str | None: """Verify the authentication token.""" - serializer = URLSafeTimedSerializer(flask.current_app.config["SECRET_KEY"]) - try: - username = serializer.loads(token, salt="auth", max_age=max_age) - except Exception: - return None - - assert isinstance(username, str) - return username - - -def generate_auth_token(username: str) -> str: - """Generate a secure authentication token.""" - serializer = URLSafeTimedSerializer(flask.current_app.config["SECRET_KEY"]) - token = typing.cast(str, serializer.dumps(username, salt="auth")) - assert isinstance(token, str) - return token + return verify_secure_token(token, "auth", max_age) def require_authentication() -> werkzeug.Response | None: @@ -45,9 +57,33 @@ def require_authentication() -> werkzeug.Response | None: if token and verify_auth_token(token): return None + callback_url = flask.url_for("auth_callback", _external=True) + + token_payload = {"original_url": flask.request.url, "callback_url": callback_url} + token = generate_secure_token(json.dumps(token_payload), salt="secure-redirect") + # Construct the redirect URL with the original URL as a parameter - return flask.redirect( - flask.current_app.config["UNIAUTH_URL"] - + "/login?next=" - + werkzeug.urls.url_quote(flask.request.url) + redirect_url = flask.current_app.config["UNIAUTH_URL"] + "/login" + redirect_to_uniauth = redirect_url + "?token=" + token + return flask.redirect(redirect_to_uniauth) + + +def auth_callback() -> tuple[str, int] | werkzeug.Response: + """Process the authentication callback.""" + auth_token = flask.request.args.get("auth_token") + token = flask.request.args["next"] # The original token passed to UniAuth + auth_token = flask.request.args[ + "auth_token" + ] # The original token passed to UniAuth + expire_date = datetime.now() + timedelta(days=180) + + original_url = verify_secure_token(token, salt="secure-redirect", max_age=600) + if not original_url: + return "Invalid or expired token", 400 + # Proceed with setting the auth_token cookie and redirecting to the original_url + # This is where you set the auth_token received from UniAuth in the client's cookies + response = flask.make_response(flask.redirect(original_url)) + response.set_cookie( + "auth_token", auth_token, expires=expire_date, httponly=True, secure=True ) + return response diff --git a/main.py b/main.py index 386fa45..e771030 100755 --- a/main.py +++ b/main.py @@ -2,6 +2,7 @@ """Single sign-on application.""" import functools +import json import typing from datetime import datetime, timedelta from typing import Any, Callable @@ -55,7 +56,7 @@ def check_user_auth() -> dict[str, Any] | None: @app.route("/login", methods=["GET", "POST"]) -def login_page() -> str | werkzeug.Response: +def login_page() -> str | werkzeug.Response | tuple[str, int]: """Login page.""" if flask.request.method == "GET": return flask.render_template("login.html") @@ -64,18 +65,28 @@ def login_page() -> str | werkzeug.Response: # 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") + 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) - flask.flash("Welcome back! You have successfully logged in.") - token = UniAuth.auth.generate_auth_token(user["username"]) - - response = flask.redirect(redirect_to) + response = flask.redirect(redirect_to_callback) response.set_cookie( "auth_token", token,