FreeBug/backend/utils/auth.py

from functools import wraps
from flask import request, jsonify, g
from sqlalchemy import select, and_
from db.model import User, Session, db
from constants import UserRole

def auth_required():
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            auth_header = request.headers.get('Authorization')
            if not auth_header:
                return jsonify({'error': 'No authorization header sent'}), 401
            try:
                session_key = auth_header.split(' ')[1]
            except IndexError:
                return jsonify({'error': 'Invalid authorization header format'}), 401
            session: Session = db.session.execute(
                select(Session).where(and_(Session.key == session_key, Session.isValid == True))
            ).scalar()
            if not session:
                return jsonify({'error': 'Invalid or expired session'}), 401
            user: User = session.user
            if not user:
                return jsonify({'error': 'User not found for the Access token'}), 401
            g.current_session = session
            g.current_user = user
            g.is_authed = True
            return f(*args, **kwargs)
        return decorated_function
    return decorator

def requires_role(roles=None):
    if roles is None:
        roles = [UserRole.USER, UserRole.ADMIN]
    roles = [int(r) for r in roles]
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if g.get('is_authed', False) is False:
                return jsonify({'error': 'Unauthorized'}), 401
            if g.current_user.role in roles:
                return f(*args, **kwargs)
            else:
                return jsonify({'error': 'Forbidden'}), 403
        return decorated_function
    return decorator