Rename cookie and refactor code

This commit is contained in:
Edward Betts 2024-02-19 09:45:17 +00:00
parent 89b111d40c
commit 0ea62afd1f

59
main.py
View file

@ -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,