More clever auth with tokens

Closes: #1

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

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,