Improvde session handling & partial code cleanup

app/__init__.py

@@ -8,6 +8,9 @@
 from flask import Flask
 from flask_wtf.csrf import CSRFProtect
 
+from flask import session
+from app.models.user import is_session_valid
+
 from app.services.database import init_db
 from app.services.session import init_session
 from app.services.recaptcha import init_recaptcha
@@ -105,6 +108,19 @@ def inject_globals():
             'PLATFORM_CHOICES': PLATFORM_CHOICES,
         }
 
+    @app.before_request
+    def validate_session():
+        """Check the session before every request is handled"""
+        name = session.get("name")
+        sess_id = session.get("sess_id")
+
+        if not name or not sess_id:
+            return
+
+        # clear invalid session
+        if not is_session_valid(name, sess_id):
+            session.clear()
+
     # Register template filters
     register_filters(app)
 

app/controllers/db_decorators.py

@@ -0,0 +1,18 @@
+"""
+Decorators for the database.
+"""
+
+from functools import wraps
+from app.models.errors import ErrUnavailable
+from app.services.database import check_connection
+
+def check_db_con(f):
+    """Decorator to check whether the database is available"""
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        if not check_connection():
+            raise ErrUnavailable("Database is unavailable")
+
+        return f(*args, **kwargs)
+
+    return wrapper

app/controllers/decorators.py

@@ -3,8 +3,8 @@
 """
 
 from functools import wraps
-from flask import redirect, url_for, session
-
+from flask import redirect, session
+#from app.models.user import is_session_valid
 
 def login_required(f):
     """Decorator to require user to be logged in."""
@@ -20,7 +20,7 @@ def anonymous_required(f):
     """Decorator to require user to NOT be logged in."""
     @wraps(f)
     def decorated_function(*args, **kwargs):
-        if session.get('name') is not None:
+        if session.get('name') is not None: #or is_session_valid(session.get('name'), session.get('sess_id')):
             return redirect('/')
         return f(*args, **kwargs)
     return decorated_function

app/controllers/login.py

@@ -5,12 +5,13 @@
 from urllib.parse import urlparse
 
 from flask import Blueprint, render_template, request, redirect, flash, session
-from app.models.user import user_by_name, user_by_mail
+from app.models.user import user_by_name, user_by_mail, add_user_session, remove_user_session
 from app.models.errors import ErrNoResult
 from app.services.passhash import match_string
 from app.services.limiter import limit
 from app.services.view import FLASH_ERROR, FLASH_SUCCESS, FLASH_WARNING, authorized_chars_only
 from app.controllers.decorators import anonymous_required
+from uuid import uuid4 # session id generator
 
 login_bp = Blueprint('login', __name__)
 
@@ -85,6 +86,11 @@ def login_post():
         if match_string(user['password'], password):
             # Login successful
             clear_main_auth()
+
+            sess_id = str(uuid4())
+            add_user_session(user['name'], sess_id)
+
+            session["sess_id"] = sess_id
             session['email'] = user['email']
             session['name'] = user['name']
             flash('Login successful!', FLASH_SUCCESS)
@@ -111,6 +117,7 @@ def login_post():
 def logout():
     """Log out the user."""
     if session.get('name'):
+        remove_user_session(session.get('name'), session.get('sess_id'))
         clear_main_auth()
         flash('Goodbye!', 'alert-info')
 

app/controllers/password.py

@@ -3,7 +3,7 @@
 """
 
 from flask import Blueprint, render_template, request, session
-from app.models.user import user_by_name, update_user_password
+from app.models.user import user_by_name, update_user_password, purge_user_sessions
 from app.services.passhash import hash_string, match_string
 from app.services.limiter import limit
 from app.controllers.decorators import login_required
@@ -65,4 +65,6 @@ def change_password_post():
         print(f"Error updating user password: {e}")
         return 'Could not update password', 500
 
+    # delete all sessions except the current one
+    purge_user_sessions(session.get('name'), [session.get('session')])
     return 'Password has been successfully updated', 200

app/models/user.py

@@ -3,22 +3,20 @@
 """
 
 import re
-from datetime import datetime
 from bson import ObjectId
 from app.services.database import get_collection, check_connection
-from app.models.errors import ErrNoResult, ErrUnavailable
-
+from app.models.errors import ErrNoResult
+from app.controllers.db_decorators import check_db_con
 
+@check_db_con
 def count_users():
     """Return the total number of users."""
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     return collection.count_documents({})
 
-
-def user_by_name(name):
+@check_db_con
+def user_by_name(name: str):
     """Get user by username (case-insensitive).
 
     Args:
@@ -31,8 +29,6 @@ def user_by_name(name):
         ErrNoResult: If user not found
         ErrUnavailable: If database unavailable
     """
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     # Case-insensitive regex search
@@ -44,8 +40,8 @@ def user_by_name(name):
 
     return result
 
-
-def user_by_mail(email):
+@check_db_con
+def user_by_mail(email: str):
     """Get user by email (case-insensitive).
 
     Args:
@@ -58,8 +54,6 @@ def user_by_mail(email):
         ErrNoResult: If user not found
         ErrUnavailable: If database unavailable
     """
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     pattern = re.compile(f'^{re.escape(email)}$', re.IGNORECASE)
@@ -70,8 +64,9 @@ def user_by_mail(email):
 
     return result
 
-
-def user_by_hexid(hexid):
+#NOTE: unused function
+@check_db_con
+def user_by_hexid(hexid: str):
     """Get user by hex ID.
 
     Args:
@@ -84,8 +79,6 @@ def user_by_hexid(hexid):
         ErrNoResult: If user not found
         ErrUnavailable: If database unavailable
     """
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     result = collection.find_one({'hexid': hexid})
@@ -95,24 +88,21 @@ def user_by_hexid(hexid):
 
     return result
 
-
+#NOTE: unused function
+@check_db_con
 def all_users_visible():
     """Get all visible users.
 
     Returns:
         List of user documents
     """
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     return list(collection.find({'visible': True}))
 
-
-def users_by_hexids(hexids):
+@check_db_con
+def users_by_hexids(hexids: list[str] | set[str]):
     """Return visible users keyed by immutable hexid."""
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     hexids = list(hexids)
     if not hexids:
@@ -129,8 +119,8 @@ def users_by_hexids(hexids):
         })
     }
 
-
-def user_create(name, email, password):
+@check_db_con
+def user_create(name: str, email: str, password: str):
     """Create a new user.
 
     Args:
@@ -141,8 +131,6 @@ def user_create(name, email, password):
     Raises:
         ErrUnavailable: If database unavailable
     """
-    if not check_connection():
-        raise ErrUnavailable("Database is unavailable")
 
     collection = get_collection('user')
     obj_id = ObjectId()
@@ -159,8 +147,8 @@ def user_create(name, email, password):
 
     collection.insert_one(user)
 
-
-def update_user_password(username, hashed_password):
+@check_db_con
+def update_user_password(username: str, hashed_password: str):
     """Update user password.
 
     Args:
@@ -186,8 +174,8 @@ def update_user_password(username, hashed_password):
     if result.matched_count == 0:
         raise ErrNoResult("No user found with the provided username")
 
-
-def user_get_unread_notifications(username):
+#don't use @check_db_con because this function returns a value
+def user_get_unread_notifications(username: str):
     """Get the count of unread notifications for a user.
 
     Args:
@@ -208,32 +196,28 @@ def user_get_unread_notifications(username):
     except Exception:
         return 0
 
-
-def user_increment_unread_notifications(username):
+@check_db_con
+def user_increment_unread_notifications(username: str):
     """Increment the unread notification count for a user.
 
     Args:
         username: The username (exact match)
     """
-    if not check_connection():
-        return
 
     collection = get_collection('user')
     collection.update_one(
         {'name': username},
         {'$inc': {'unread_notifications': 1}}
     )
 
-
-def user_decrement_unread_notifications(username, count=1):
+@check_db_con
+def user_decrement_unread_notifications(username: str, count: int = 1):
     """Decrement the unread notification count for a user.
 
     Args:
         username: The username (exact match)
         count: Number to decrement by (default 1)
     """
-    if not check_connection():
-        return
 
     collection = get_collection('user')
 
@@ -246,3 +230,87 @@ def user_decrement_unread_notifications(username, count=1):
             {'name': username},
             {'$set': {'unread_notifications': new_count}}
         )
+
+@check_db_con
+def is_session_valid(username: str | None, sess_id: str | None) -> bool:
+    """Check whether a session is valid
+
+    Args:
+        username: The username, used to bind user and session
+        sess_id: The UUIDv4 session key as str or None
+
+    Returns:
+        False if sess_id is None, or sess_id isn't valid
+        True if sess_id is valid for username 
+    """
+    if not sess_id or not username:
+        return False
+
+    collection = get_collection('user')
+
+    pattern = re.compile(f'^{re.escape(username)}$', re.IGNORECASE)
+    target_user = collection.find_one({'name': pattern})
+
+    if not target_user:
+        return False
+
+    if not "active_sessions" in target_user: # expected for not migrated accounts
+        collection.update_one(
+            {"_id": target_user["_id"]},
+            {"$set": {"active_sessions": []}}
+        )
+        # non-migrated accounts should be logged out
+        return False
+
+    return sess_id in target_user["active_sessions"]
+
+@check_db_con
+def add_user_session(username: str, sess_id: str) -> None:
+    """Add an active session key for a user
+    Args:
+        username: The username to add the session key to
+        sess_id: The UUIDv4 session key
+    """
+    collection = get_collection('user')
+
+    pattern = re.compile(f'^{re.escape(username)}$', re.IGNORECASE)
+    target_user = collection.find_one({'name': pattern})
+
+    collection.update_one(
+        {"_id": target_user["_id"]},
+        {"$addToSet": {"active_sessions": sess_id}}
+    )
+
+@check_db_con
+def remove_user_session(username: str, sess_id: str) -> None:
+    """Remove an active session key from a user
+    Args:
+        username: The target username
+        sess_id: The session id to invalidate/delete
+    """
+    collection = get_collection('user')
+
+    pattern = re.compile(f'^{re.escape(username)}$', re.IGNORECASE)
+    target_user = collection.find_one({'name': pattern})
+
+    collection.update_one(
+        {"_id": target_user["_id"]},
+        {"$pull": {"active_sessions": sess_id}}
+    )
+
+@check_db_con
+def purge_user_sessions(username: str, exclude: list[str] = []) -> None:
+    """Delete all sessions (exlcuding the ones in the exclude list)
+    Args:
+        username: Target username
+        exclude: A list of sessions to NOT delete, leave empty to delete ALL sessions
+    """
+    collection = get_collection('user')
+
+    pattern = re.compile(f'^{re.escape(username)}$', re.IGNORECASE)
+    target_user = collection.find_one({'name': pattern})
+
+    collection.update_one(
+        {"_id": target_user["_id"]},
+        {"$set": {"active_sessions": exclude}}
+    )