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)) diff --git a/main.py b/main.py index e771030..ff37417 100755 --- a/main.py +++ b/main.py @@ -13,7 +13,7 @@ import werkzeug import UniAuth.auth app = flask.Flask(__name__) -app.debug = True +app.debug = False app.config.from_object("config.default") @@ -22,7 +22,7 @@ def login_required(f: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(f) def decorated_function(*args, **kwargs) -> werkzeug.Response: # type: ignore - token = flask.request.cookies.get("auth_token") + token = flask.request.cookies.get("uniauth_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 @@ -55,40 +55,67 @@ def check_user_auth() -> dict[str, Any] | 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") - +def read_callback_url_from_token() -> str | None: + """Parse token and extract callback URL.""" token = flask.request.args.get("token") if not token: - return "Invalid token", 400 + return None json_data = UniAuth.auth.verify_secure_token( token, salt="secure-redirect", max_age=600 ) + if not json_data: + return None 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"] + assert isinstance(callback_url, str) + + return callback_url + + +@app.route("/login", methods=["GET", "POST"]) +def login_page() -> str | werkzeug.Response | tuple[str, int]: + """Login page.""" + app.logger.info("Login page.") + if flask.request.method == "GET": + uniauth_token = flask.request.cookies.get("uniauth_token") + if ( + not uniauth_token + or not UniAuth.auth.verify_auth_token(uniauth_token) + or not (callback_url := read_callback_url_from_token()) + ): + return flask.render_template("login.html") + token = flask.request.args["token"] + redirect_to_callback = f"{callback_url}?auth_token={uniauth_token}&next={token}" + app.logger.info(f"Redirecting to: {redirect_to_callback}") + + return flask.redirect(redirect_to_callback) + + if not (user := check_user_auth()): + # Login failed: Show an error message on the login page + app.logger.info("User auth failed") + return flask.render_template("login.html", error="Invalid credentials") + + callback_url = read_callback_url_from_token() + if not callback_url: + return "Invalid token", 400 + auth_token = UniAuth.auth.generate_auth_token(user["username"]) + token = flask.request.args["token"] redirect_to_callback = f"{callback_url}?auth_token={auth_token}&next={token}" + app.logger.info(f"Redirecting to: {redirect_to_callback}") + # 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", + "uniauth_token", token, expires=expire_date, httponly=True,