Move authentication decorators to common utils.py

- Created vidai/utils.py with shared authentication decorators
- Removed duplicate decorator definitions from web.py, api.py, and admin.py
- Updated all modules to import decorators from utils.py
- Maintains consistent authentication logic across all modules

vidai/admin.py

@@ -23,6 +23,7 @@ from flask import Blueprint, request, render_template, redirect, url_for, flash
 from .auth import require_auth, require_admin
 from .database import get_user_tokens, update_user_tokens, get_user_queue_items, get_default_user_tokens, create_remember_token, validate_remember_token, delete_remember_token, extend_remember_token, get_all_users, update_user_status, update_user_info, delete_user, get_worker_tokens, deactivate_worker_token, activate_worker_token, delete_worker_token, create_user
 from .comm import SocketCommunicator, Message
+from .utils import get_current_user_session, login_required, admin_required
 
 admin_bp = Blueprint('admin', __name__, url_prefix='/admin')
 
@@ -66,14 +67,6 @@ def get_result(msg_id: str) -> dict:
             time.sleep(0.1)
     return {'error': 'Timeout waiting for result'}
 
-def get_current_user_session():
-    """Get current user from session."""
-    from flask import session
-    from .auth import get_current_user
-    session_id = session.get('session_id')
-    if session_id:
-        return get_current_user(session_id)
-    return None
 
 @admin_bp.route('/train', methods=['GET', 'POST'])
 @require_admin

vidai/api.py

@@ -25,6 +25,7 @@ import uuid
 from .auth import login_user, logout_user, get_current_user, register_user, confirm_email, require_auth, require_admin
 from .database import get_user_tokens, update_user_tokens, get_user_queue_items, get_default_user_tokens, create_remember_token, validate_remember_token, delete_remember_token, extend_remember_token
 from .comm import SocketCommunicator, Message
+from .utils import get_current_user_session, login_required, admin_required, api_auth_required, admin_api_auth_required
 
 api_bp = Blueprint('api', __name__)
 
@@ -34,93 +35,6 @@ server_dir = None
 # Communicator to backend (always TCP)
 comm = SocketCommunicator(host='localhost', port=5001, comm_type='tcp')
 
-def get_current_user_session():
-    """Get current user from session."""
-    from flask import session
-    session_id = session.get('session_id')
-    if session_id:
-        return get_current_user(session_id)
-    return None
-
-def login_required(f):
-    """Decorator to require login."""
-    def decorated_function(*args, **kwargs):
-        user = get_current_user_session()
-        if not user:
-            from flask import redirect, url_for
-            return redirect(url_for('login'))
-        return f(*args, **kwargs)
-    decorated_function.__name__ = f.__name__
-    return decorated_function
-
-def admin_required(f):
-    """Decorator to require admin role."""
-    def decorated_function(*args, **kwargs):
-        user = get_current_user_session()
-        if not user or user['role'] != 'admin':
-            from flask import flash, redirect, url_for
-            flash('Admin access required', 'error')
-            return redirect(url_for('dashboard'))
-        return f(*args, **kwargs)
-    decorated_function.__name__ = f.__name__
-    return decorated_function
-
-def api_auth_required(f):
-    """Decorator to require authentication via session or API token."""
-    def decorated_function(*args, **kwargs):
-        # Check for session auth first
-        user = get_current_user_session()
-        if user:
-            request.api_user = user
-            return f(*args, **kwargs)
-
-        # Check for API token auth
-        auth_header = request.headers.get('Authorization')
-        if auth_header and auth_header.startswith('Bearer '):
-            token = auth_header.split(' ')[1]
-            from .database import validate_user_api_token
-            user = validate_user_api_token(token)
-            if user:
-                request.api_user = user
-                return f(*args, **kwargs)
-
-        # No valid auth
-        if request.is_json or request.path.startswith('/api/'):
-            return json.dumps({'error': 'Authentication required'}), 401
-        else:
-            from flask import redirect, url_for
-            return redirect(url_for('login'))
-    decorated_function.__name__ = f.__name__
-    return decorated_function
-
-def admin_api_auth_required(f):
-    """Decorator to require admin authentication via session or API token."""
-    def decorated_function(*args, **kwargs):
-        # Check for session auth first
-        user = get_current_user_session()
-        if user and user.get('role') == 'admin':
-            request.api_user = user
-            return f(*args, **kwargs)
-
-        # Check for API token auth
-        auth_header = request.headers.get('Authorization')
-        if auth_header and auth_header.startswith('Bearer '):
-            token = auth_header.split(' ')[1]
-            from .database import validate_user_api_token
-            user = validate_user_api_token(token)
-            if user and user.get('role') == 'admin':
-                request.api_user = user
-                return f(*args, **kwargs)
-
-        # No valid auth
-        if request.is_json or request.path.startswith('/api/'):
-            return json.dumps({'error': 'Admin authentication required'}), 401
-        else:
-            from flask import redirect, url_for, flash
-            flash('Admin access required', 'error')
-            return redirect(url_for('dashboard'))
-    decorated_function.__name__ = f.__name__
-    return decorated_function
 
 def send_to_backend(msg_type: str, data: dict) -> str:
     """Send message to backend and return message id."""

vidai/utils.py

+# Video AI Utilities
+# Copyright (C) 2024 Stefy Lanza <stefy@sexhack.me>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""
+Common utilities for Video AI application.
+Includes authentication decorators and helper functions.
+"""
+
+from .auth import get_current_user
+
+def get_current_user_session():
+    """Get current user from session."""
+    from flask import session
+    session_id = session.get('session_id')
+    if session_id:
+        return get_current_user(session_id)
+    return None
+
+def login_required(f):
+    """Decorator to require login."""
+    def decorated_function(*args, **kwargs):
+        user = get_current_user_session()
+        if not user:
+            from flask import redirect, url_for
+            return redirect(url_for('login'))
+        return f(*args, **kwargs)
+    decorated_function.__name__ = f.__name__
+    return decorated_function
+
+def admin_required(f):
+    """Decorator to require admin role."""
+    def decorated_function(*args, **kwargs):
+        user = get_current_user_session()
+        if not user or user['role'] != 'admin':
+            from flask import flash, redirect, url_for
+            flash('Admin access required', 'error')
+            return redirect(url_for('dashboard'))
+        return f(*args, **kwargs)
+    decorated_function.__name__ = f.__name__
+    return decorated_function
+
+def api_auth_required(f):
+    """Decorator to require authentication via session or API token."""
+    def decorated_function(*args, **kwargs):
+        # Check for session auth first
+        user = get_current_user_session()
+        if user:
+            from flask import request
+            request.api_user = user
+            return f(*args, **kwargs)
+
+        # Check for API token auth
+        from flask import request
+        auth_header = request.headers.get('Authorization')
+        if auth_header and auth_header.startswith('Bearer '):
+            token = auth_header.split(' ')[1]
+            from .database import validate_user_api_token
+            user = validate_user_api_token(token)
+            if user:
+                request.api_user = user
+                return f(*args, **kwargs)
+
+        # No valid auth
+        if request.is_json or request.path.startswith('/api/'):
+            return {'error': 'Authentication required'}, 401
+        else:
+            from flask import redirect, url_for
+            return redirect(url_for('login'))
+    decorated_function.__name__ = f.__name__
+    return decorated_function
+
+def admin_api_auth_required(f):
+    """Decorator to require admin authentication via session or API token."""
+    def decorated_function(*args, **kwargs):
+        # Check for session auth first
+        user = get_current_user_session()
+        if user and user.get('role') == 'admin':
+            from flask import request
+            request.api_user = user
+            return f(*args, **kwargs)
+
+        # Check for API token auth
+        from flask import request
+        auth_header = request.headers.get('Authorization')
+        if auth_header and auth_header.startswith('Bearer '):
+            token = auth_header.split(' ')[1]
+            from .database import validate_user_api_token
+            user = validate_user_api_token(token)
+            if user and user.get('role') == 'admin':
+                request.api_user = user
+                return f(*args, **kwargs)
+
+        # No valid auth
+        if request.is_json or request.path.startswith('/api/'):
+            return {'error': 'Admin authentication required'}, 401
+        else:
+            from flask import redirect, url_for, flash
+            flash('Admin access required', 'error')
+            return redirect(url_for('dashboard'))
+    decorated_function.__name__ = f.__name__
+    return decorated_function
\ No newline at end of file

vidai/web.py

@@ -31,6 +31,7 @@ from .auth import login_user, logout_user, get_current_user, register_user, conf
 from .database import get_user_tokens, update_user_tokens, get_user_queue_items, get_default_user_tokens, create_remember_token, validate_remember_token, delete_remember_token, extend_remember_token
 from .api import api_bp
 from .admin import admin_bp
+from .utils import get_current_user_session, login_required, admin_required
 
 app = Flask(__name__, template_folder='../templates')
 app.secret_key = os.environ.get('FLASK_SECRET_KEY', 'dev-secret-key-change-in-production')
@@ -55,33 +56,6 @@ def check_remember_me():
 # Communicator to backend (always TCP)
 comm = SocketCommunicator(host='localhost', port=5001, comm_type='tcp')
 
-def get_current_user_session():
-    """Get current user from session."""
-    session_id = session.get('session_id')
-    if session_id:
-        return get_current_user(session_id)
-    return None
-
-def login_required(f):
-    """Decorator to require login."""
-    def decorated_function(*args, **kwargs):
-        user = get_current_user_session()
-        if not user:
-            return redirect(url_for('login'))
-        return f(*args, **kwargs)
-    decorated_function.__name__ = f.__name__
-    return decorated_function
-
-def admin_required(f):
-    """Decorator to require admin role."""
-    def decorated_function(*args, **kwargs):
-        user = get_current_user_session()
-        if not user or user['role'] != 'admin':
-            flash('Admin access required', 'error')
-            return redirect(url_for('dashboard'))
-        return f(*args, **kwargs)
-    decorated_function.__name__ = f.__name__
-    return decorated_function
 
 def send_to_backend(msg_type: str, data: dict) -> str:
     """Send message to backend and return message id."""