diff --git a/UniAuth/auth.py b/UniAuth/auth.py index 20f8446..75b6deb 100644 --- a/UniAuth/auth.py +++ b/UniAuth/auth.py @@ -5,6 +5,7 @@ import typing from datetime import datetime, timedelta import flask +import itsdangerous import werkzeug from itsdangerous.url_safe import URLSafeTimedSerializer @@ -29,7 +30,7 @@ def verify_secure_token(token: str, salt: str, max_age: int) -> str | None: serializer = URLSafeTimedSerializer(flask.current_app.config["SECRET_KEY"]) try: data = serializer.loads(token, salt=salt, max_age=max_age) - except Exception: + except itsdangerous.exc.BadTimeSignature: return None assert isinstance(data, str) @@ -53,6 +54,9 @@ def require_authentication() -> werkzeug.Response | None: if not flask.current_app.config.get("REQUIRE_AUTH"): return None + if flask.request.endpoint == "auth_callback": + return None + token = flask.request.cookies.get("auth_token") if token and verify_auth_token(token): return None @@ -77,9 +81,11 @@ def auth_callback() -> tuple[str, int] | werkzeug.Response: ] # 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: + token_payload = verify_secure_token(token, salt="secure-redirect", max_age=600) + if not token_payload: return "Invalid or expired token", 400 + original_url = json.loads(token_payload)["original_url"] + # 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))