|
|
from functools import wraps |
|
|
|
|
|
from flask import current_app, g, has_request_context, request |
|
|
from flask_login import user_logged_in |
|
|
from flask_login.config import EXEMPT_METHODS |
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
current_user = LocalProxy(lambda: _get_user()) |
|
|
|
|
|
|
|
|
def login_required(func): |
|
|
""" |
|
|
If you decorate a view with this, it will ensure that the current user is |
|
|
logged in and authenticated before calling the actual view. (If they are |
|
|
not, it calls the :attr:`LoginManager.unauthorized` callback.) For |
|
|
example:: |
|
|
|
|
|
@app.route('/post') |
|
|
@login_required |
|
|
def post(): |
|
|
pass |
|
|
|
|
|
If there are only certain times you need to require that your user is |
|
|
logged in, you can do so with:: |
|
|
|
|
|
if not current_user.is_authenticated: |
|
|
return current_app.login_manager.unauthorized() |
|
|
|
|
|
...which is essentially the code that this function adds to your views. |
|
|
|
|
|
It can be convenient to globally turn off authentication when unit testing. |
|
|
To enable this, if the application configuration variable `LOGIN_DISABLED` |
|
|
is set to `True`, this decorator will be ignored. |
|
|
|
|
|
.. Note :: |
|
|
|
|
|
Per `W3 guidelines for CORS preflight requests |
|
|
<http://www.w3.org/TR/cors/#cross-origin-request-with-preflight-0>`_, |
|
|
HTTP ``OPTIONS`` requests are exempt from login checks. |
|
|
|
|
|
:param func: The view function to decorate. |
|
|
:type func: function |
|
|
""" |
|
|
|
|
|
@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 = Account.query.filter_by(id=ta.account_id).first() |
|
|
|
|
|
if account: |
|
|
account.current_tenant = tenant |
|
|
current_app.login_manager._update_request_context_with_user(account) |
|
|
user_logged_in.send(current_app._get_current_object(), user=_get_user()) |
|
|
if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED: |
|
|
pass |
|
|
elif not current_user.is_authenticated: |
|
|
return current_app.login_manager.unauthorized() |
|
|
|
|
|
|
|
|
|
|
|
if callable(getattr(current_app, "ensure_sync", None)): |
|
|
return current_app.ensure_sync(func)(*args, **kwargs) |
|
|
return func(*args, **kwargs) |
|
|
|
|
|
return decorated_view |
|
|
|
|
|
|
|
|
def _get_user(): |
|
|
if has_request_context(): |
|
|
if "_login_user" not in g: |
|
|
current_app.login_manager._load_user() |
|
|
|
|
|
return g._login_user |
|
|
|
|
|
return None |
|
|
|
