dify/api/controllers/web/wraps.py

from functools import wraps

from flask import request
from flask_restful import Resource
from werkzeug.exceptions import BadRequest, NotFound, Unauthorized

from controllers.web.error import WebSSOAuthRequiredError
from extensions.ext_database import db
from libs.passport import PassportService
from models.model import App, EndUser, Site
from services.enterprise.enterprise_service import EnterpriseService
from services.feature_service import FeatureService


def validate_jwt_token(view=None):
    def decorator(view):
        @wraps(view)
        def decorated(*args, **kwargs):
            app_model, end_user = decode_jwt_token()

            return view(app_model, end_user, *args, **kwargs)

        return decorated

    if view:
        return decorator(view)
    return decorator


def decode_jwt_token():
    system_features = FeatureService.get_system_features()
    app_code = request.headers.get("X-App-Code")
    try:
        auth_header = request.headers.get("Authorization")
        if auth_header is None:
            raise Unauthorized("Authorization header is missing.")

        if " " not in auth_header:
            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")

        auth_scheme, tk = auth_header.split(None, 1)
        auth_scheme = auth_scheme.lower()

        if auth_scheme != "bearer":
            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
        decoded = PassportService().verify(tk)
        app_code = decoded.get("app_code")
        app_model = db.session.query(App).filter(App.id == decoded["app_id"]).first()
        site = db.session.query(Site).filter(Site.code == app_code).first()
        if not app_model:
            raise NotFound()
        if not app_code or not site:
            raise BadRequest("Site URL is no longer valid.")
        if app_model.enable_site is False:
            raise BadRequest("Site is disabled.")
        end_user = db.session.query(EndUser).filter(EndUser.id == decoded["end_user_id"]).first()
        if not end_user:
            raise NotFound()

        _validate_web_sso_token(decoded, system_features, app_code)

        return app_model, end_user
    except Unauthorized as e:
        if system_features.sso_enforced_for_web:
            app_web_sso_enabled = EnterpriseService.get_app_web_sso_enabled(app_code).get("enabled", False)
            if app_web_sso_enabled:
                raise WebSSOAuthRequiredError()

        raise Unauthorized(e.description)


def _validate_web_sso_token(decoded, system_features, app_code):
    app_web_sso_enabled = False

    # Check if SSO is enforced for web, and if the token source is not SSO, raise an error and redirect to SSO login
    if system_features.sso_enforced_for_web:
        app_web_sso_enabled = EnterpriseService.get_app_web_sso_enabled(app_code).get("enabled", False)
        if app_web_sso_enabled:
            source = decoded.get("token_source")
            if not source or source != "sso":
                raise WebSSOAuthRequiredError()

    # Check if SSO is not enforced for web, and if the token source is SSO,
    # raise an error and redirect to normal passport login
    if not system_features.sso_enforced_for_web or not app_web_sso_enabled:
        source = decoded.get("token_source")
        if source and source == "sso":
            raise Unauthorized("sso token expired.")


class WebApiResource(Resource):
    method_decorators = [validate_jwt_token]