Browse Source

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

Signed-off-by: -LAN- <laipz8200@outlook.com>
tags/1.4.1
-LAN- 5 months ago
parent
commit
f21e6e03a3
No account linked to committer's email address
2 changed files with 23 additions and 35 deletions
  1. 22
    1
      api/extensions/ext_login.py
  2. 1
    34
      api/libs/login.py

+ 22
- 1
api/extensions/ext_login.py View File

from werkzeug.exceptions import NotFound, Unauthorized from werkzeug.exceptions import NotFound, Unauthorized


import contexts import contexts
from configs import dify_config
from dify_app import DifyApp from dify_app import DifyApp
from extensions.ext_database import db from extensions.ext_database import db
from libs.passport import PassportService from libs.passport import PassportService
from models.account import Account
from models.account import Account, Tenant, TenantAccountJoin
from models.model import EndUser from models.model import EndUser
from services.account_service import AccountService from services.account_service import AccountService


else: else:
auth_token = request.args.get("_token") 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 request.blueprint in {"console", "inner_api"}:
if not auth_token: if not auth_token:
raise Unauthorized("Invalid Authorization token.") raise Unauthorized("Invalid Authorization token.")

+ 1
- 34
api/libs/login.py View File

from typing import Any from typing import Any


from flask import current_app, g, has_request_context, request 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 flask_login.config import EXEMPT_METHODS # type: ignore
from werkzeug.exceptions import Unauthorized
from werkzeug.local import LocalProxy from werkzeug.local import LocalProxy


from configs import dify_config 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 from models.model import EndUser


#: A proxy for the current user. If no user is logged in, this will be an #: A proxy for the current user. If no user is logged in, this will be an


@wraps(func) @wraps(func)
def decorated_view(*args, **kwargs): 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: if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED:
pass pass
elif not current_user.is_authenticated: elif not current_user.is_authenticated:

Loading…
Cancel
Save