UniAuth/main.py

103 lines
2.9 KiB
Python
Executable file

#!/usr/bin/python3
"""Single sign-on application."""
import functools
import typing
from datetime import datetime, timedelta
from typing import Any, Callable
import flask
import werkzeug
import UniAuth.auth
app = flask.Flask(__name__)
app.debug = True
app.config.from_object("config.default")
def login_required(f: Callable[..., Any]) -> Callable[..., Any]:
"""Route requires login decorator."""
@functools.wraps(f)
def decorated_function(*args, **kwargs) -> werkzeug.Response: # type: ignore
token = flask.request.cookies.get("auth_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
return flask.redirect(flask.url_for("login_page"))
return typing.cast(werkzeug.Response, f(*args, **kwargs))
return decorated_function
@app.route("/")
@login_required
def root_page() -> str:
"""Root page."""
return flask.render_template("auth_good.html")
def check_user_auth() -> dict[str, Any] | None:
"""Load username and password and check if valid."""
# Extract username and password from POST request
username = flask.request.form.get("username", type=str)
password = flask.request.form.get("password", type=str)
# Retrieve users from app config
users = app.config["USERS"]
# Check if the username and password match any user in the list
return next(
(u for u in users if u["username"] == username and u["password"] == password),
None,
)
@app.route("/login", methods=["GET", "POST"])
def login_page() -> str | werkzeug.Response:
"""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")
redirect_to = (
flask.request.args.get("next")
or flask.session.get("next")
or flask.url_for("dashboard")
)
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.set_cookie(
"auth_token",
token,
expires=expire_date,
httponly=True,
secure=True,
samesite="Lax",
)
return response
@app.route("/logout")
def logout() -> werkzeug.Response:
"""Handle user logout by clearing the authentication cookie."""
after_login = flask.request.args.get("next")
response = flask.redirect(flask.url_for("login_page", next=after_login))
response.set_cookie("auth_token", "", expires=0)
flask.flash("You have been successfully logged out.", "info")
return response
if __name__ == "__main__":
app.run(host="0.0.0.0")