diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 80fee7ccd8..6a5721c021 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -6,10 +6,11 @@ from flask_login import user_loaded_from_request, user_logged_in from werkzeug.exceptions import NotFound, Unauthorized import contexts +from configs import dify_config from dify_app import DifyApp from extensions.ext_database import db from libs.passport import PassportService -from models.account import Account +from models.account import Account, Tenant, TenantAccountJoin from models.model import EndUser from services.account_service import AccountService @@ -32,6 +33,26 @@ def load_user_from_request(request_from_flask_login): else: auth_token = request.args.get("_token") + # Check for admin API key authentication first + if dify_config.ADMIN_API_KEY_ENABLE and auth_header: + admin_api_key = dify_config.ADMIN_API_KEY + if admin_api_key and admin_api_key == auth_token: + workspace_id = request.headers.get("X-WORKSPACE-ID") + if workspace_id: + tenant_account_join = ( + db.session.query(Tenant, TenantAccountJoin) + .filter(Tenant.id == workspace_id) + .filter(TenantAccountJoin.tenant_id == Tenant.id) + .filter(TenantAccountJoin.role == "owner") + .one_or_none() + ) + if tenant_account_join: + tenant, ta = tenant_account_join + account = db.session.query(Account).filter_by(id=ta.account_id).first() + if account: + account.current_tenant = tenant + return account + if request.blueprint in {"console", "inner_api"}: if not auth_token: raise Unauthorized("Invalid Authorization token.") diff --git a/api/libs/login.py b/api/libs/login.py index be9478e850..e3a7fe2948 100644 --- a/api/libs/login.py +++ b/api/libs/login.py @@ -2,14 +2,11 @@ from functools import wraps from typing import Any from flask import current_app, g, has_request_context, request -from flask_login import user_logged_in # type: ignore from flask_login.config import EXEMPT_METHODS # type: ignore -from werkzeug.exceptions import Unauthorized from werkzeug.local import LocalProxy from configs import dify_config -from extensions.ext_database import db -from models.account import Account, Tenant, TenantAccountJoin +from models.account import Account from models.model import EndUser #: A proxy for the current user. If no user is logged in, this will be an @@ -53,36 +50,6 @@ def login_required(func): @wraps(func) def decorated_view(*args, **kwargs): - auth_header = request.headers.get("Authorization") - if dify_config.ADMIN_API_KEY_ENABLE: - if auth_header: - if " " not in auth_header: - raise Unauthorized("Invalid Authorization header format. Expected 'Bearer ' format.") - auth_scheme, auth_token = auth_header.split(None, 1) - auth_scheme = auth_scheme.lower() - if auth_scheme != "bearer": - raise Unauthorized("Invalid Authorization header format. Expected 'Bearer ' format.") - - admin_api_key = dify_config.ADMIN_API_KEY - if admin_api_key: - if admin_api_key == auth_token: - workspace_id = request.headers.get("X-WORKSPACE-ID") - if workspace_id: - tenant_account_join = ( - db.session.query(Tenant, TenantAccountJoin) - .filter(Tenant.id == workspace_id) - .filter(TenantAccountJoin.tenant_id == Tenant.id) - .filter(TenantAccountJoin.role == "owner") - .one_or_none() - ) - if tenant_account_join: - tenant, ta = tenant_account_join - account = db.session.query(Account).filter_by(id=ta.account_id).first() - # Login admin - if account: - account.current_tenant = tenant - current_app.login_manager._update_request_context_with_user(account) # type: ignore - user_logged_in.send(current_app._get_current_object(), user=_get_user()) # type: ignore if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED: pass elif not current_user.is_authenticated: