Parcourir la source

refactor: Consolidate Flask-Login Authentication Logic (#20235)

Signed-off-by: -LAN- <laipz8200@outlook.com>
tags/1.4.1
-LAN- il y a 5 mois
Parent
révision
f21e6e03a3
Aucun compte lié à l'adresse e-mail de l'auteur
2 fichiers modifiés avec 23 ajouts et 35 suppressions
  1. 22
    1
      api/extensions/ext_login.py
  2. 1
    34
      api/libs/login.py

+ 22
- 1
api/extensions/ext_login.py Voir le fichier

@@ -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.")

+ 1
- 34
api/libs/login.py Voir le fichier

@@ -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 <api-key>' 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 <api-key>' 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:

Chargement…
Annuler
Enregistrer