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,