From 5d303145310cc7da613f7a85d8fd3d5643ee25ec Mon Sep 17 00:00:00 2001
From: Edward Betts <edward@4angle.com>
Date: Tue, 23 Jan 2024 10:44:16 +0000
Subject: [PATCH] New UniAuth.auth to be used by other services

---
 UniAuth/__init__.py |  0
 UniAuth/auth.py     | 53 +++++++++++++++++++++++++++++++++++++++++++++
 main.py             | 30 +++++--------------------
 3 files changed, 59 insertions(+), 24 deletions(-)
 create mode 100644 UniAuth/__init__.py
 create mode 100644 UniAuth/auth.py

diff --git a/UniAuth/__init__.py b/UniAuth/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/UniAuth/auth.py b/UniAuth/auth.py
new file mode 100644
index 0000000..62c3a34
--- /dev/null
+++ b/UniAuth/auth.py
@@ -0,0 +1,53 @@
+"""Authentication via UniAuth."""
+
+import typing
+
+import flask
+import werkzeug
+from itsdangerous.url_safe import URLSafeTimedSerializer
+
+max_age = 60 * 60 * 24 * 90
+
+
+def is_authentication() -> bool:
+    """User is authenticated."""
+    return not flask.current_app.config.get("REQUIRE_AUTH") or bool(
+        (token := flask.request.cookies.get("auth_token")) and verify_auth_token(token)
+    )
+
+
+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
+
+
+def require_authentication() -> werkzeug.Response | None:
+    """Require authentication and redirect with return URL."""
+    if not flask.current_app.config.get("REQUIRE_AUTH"):
+        return None
+
+    token = flask.request.cookies.get("auth_token")
+    if token and verify_auth_token(token):
+        return None
+
+    # 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)
+    )
diff --git a/main.py b/main.py
index c0a8bb7..386fa45 100755
--- a/main.py
+++ b/main.py
@@ -8,33 +8,13 @@ from typing import Any, Callable
 
 import flask
 import werkzeug
-from itsdangerous.url_safe import URLSafeTimedSerializer
+
+import UniAuth.auth
 
 app = flask.Flask(__name__)
 app.debug = True
 app.config.from_object("config.default")
 
-serializer = URLSafeTimedSerializer(app.config["SECRET_KEY"])
-max_age = 60 * 60 * 24 * 90
-
-
-def generate_auth_token(username: str) -> str:
-    """Generate a secure authentication token."""
-    token = typing.cast(str, serializer.dumps(username, salt="auth"))
-    assert isinstance(token, str)
-    return token
-
-
-def verify_auth_token(token: str) -> str | None:
-    """Verify the authentication token."""
-    try:
-        username = serializer.loads(token, salt="auth", max_age=max_age)
-    except Exception:
-        return None
-
-    assert isinstance(username, str)
-    return username
-
 
 def login_required(f: Callable[..., Any]) -> Callable[..., Any]:
     """Route requires login decorator."""
@@ -42,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")
-        if not token or verify_auth_token(token) is None:
+        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"))
@@ -93,10 +73,12 @@ def login_page() -> str | werkzeug.Response:
     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",
-        generate_auth_token(user["username"]),
+        token,
         expires=expire_date,
         httponly=True,
         secure=True,