From 3757cbdb21aa7b105c5d6b4ddbbbf88aa950e5bf Mon Sep 17 00:00:00 2001
From: Sam Chau <chaucted@gmail.com>
Date: Sun, 2 Feb 2025 09:08:20 +1030
Subject: [PATCH] feat: add session ID management for authentication setup and
 middleware

---
 backend/app/auth/__init__.py                  | 20 ++++++++++++++++---
 .../migrations/versions/001_initial_schema.py |  1 +
 backend/app/middleware.py                     |  7 ++++++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/backend/app/auth/__init__.py b/backend/app/auth/__init__.py
index 41cd712..77d28c3 100644
--- a/backend/app/auth/__init__.py
+++ b/backend/app/auth/__init__.py
@@ -1,3 +1,5 @@
+# backend/app/auth/__init__.py
+
 from flask import Blueprint, jsonify, request, session
 from werkzeug.security import generate_password_hash, check_password_hash
 import secrets
@@ -19,7 +21,6 @@ def setup():
         return jsonify({'needs_setup': True}), 200
 
     # Handle POST request for actual setup
-    # Check if auth already exists
     if db.execute('SELECT 1 FROM auth').fetchone():
         logger.warning('Failed setup attempt - auth already configured')
         return jsonify({'error': 'Auth already configured'}), 400
@@ -34,16 +35,18 @@ def setup():
 
     api_key = secrets.token_urlsafe(32)
     password_hash = generate_password_hash(password)
+    session_id = secrets.token_urlsafe(32)  # Generate a new session ID
 
     try:
         db.execute(
-            'INSERT INTO auth (username, password_hash, api_key) VALUES (?, ?, ?)',
-            (username, password_hash, api_key))
+            'INSERT INTO auth (username, password_hash, api_key, session_id) VALUES (?, ?, ?, ?)',
+            (username, password_hash, api_key, session_id))
         db.commit()
         logger.info('Initial auth setup completed successfully')
 
         # Set up session after successful creation
         session['authenticated'] = True
+        session['session_id'] = session_id
         session.permanent = True
 
         return jsonify({
@@ -86,12 +89,23 @@ def authenticate():
                       (username, )).fetchone()
 
     if user and check_password_hash(user['password_hash'], password):
+        # Generate a new session ID
+        new_session_id = secrets.token_urlsafe(32)
+        db.execute('UPDATE auth SET session_id = ? WHERE username = ?',
+                   (new_session_id, username))
+        db.commit()
+
+        # Set up session
         session['authenticated'] = True
+        session[
+            'session_id'] = new_session_id  # Store session ID in the session
         session.permanent = True
+
         # Clear failed attempts on success
         db.execute('DELETE FROM failed_attempts WHERE ip_address = ?',
                    (ip_address, ))
         db.commit()
+
         logger.info(f'Successful authentication for user: {username}')
         return jsonify({'authenticated': True})
 
diff --git a/backend/app/db/migrations/versions/001_initial_schema.py b/backend/app/db/migrations/versions/001_initial_schema.py
index c0ccadf..e764786 100644
--- a/backend/app/db/migrations/versions/001_initial_schema.py
+++ b/backend/app/db/migrations/versions/001_initial_schema.py
@@ -70,6 +70,7 @@ def up():
             username TEXT NOT NULL,
             password_hash TEXT NOT NULL,
             api_key TEXT,
+            session_id TEXT, 
             created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
         )
         ''')
diff --git a/backend/app/middleware.py b/backend/app/middleware.py
index fe2f4b8..bb4e8e3 100644
--- a/backend/app/middleware.py
+++ b/backend/app/middleware.py
@@ -1,3 +1,5 @@
+# backend/app/middleware.py
+
 from functools import wraps
 from flask import request, session, jsonify, current_app
 from .db import get_db
@@ -27,7 +29,10 @@ def init_middleware(app):
 
         # Check session authentication (for web users)
         if session.get('authenticated'):
-            return
+            db = get_db()
+            user = db.execute('SELECT session_id FROM auth').fetchone()
+            if user and session.get('session_id') == user['session_id']:
+                return
 
         # Check API key authentication (for API users)
         api_key = request.headers.get('X-Api-Key')