dify

dify/api/controllers/console/workspace/__init__.py

@@ -0,0 +1,62 @@
+from collections.abc import Callable
+from functools import wraps
+from typing import ParamSpec, TypeVar
+
+from sqlalchemy.orm import Session
+from werkzeug.exceptions import Forbidden
+
+from extensions.ext_database import db
+from libs.login import current_account_with_tenant
+from models.account import TenantPluginPermission
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def plugin_permission_required(
+    install_required: bool = False,
+    debug_required: bool = False,
+):
+    def interceptor(view: Callable[P, R]):
+        @wraps(view)
+        def decorated(*args: P.args, **kwargs: P.kwargs):
+            current_user, current_tenant_id = current_account_with_tenant()
+            user = current_user
+            tenant_id = current_tenant_id
+
+            with Session(db.engine) as session:
+                permission = (
+                    session.query(TenantPluginPermission)
+                    .where(
+                        TenantPluginPermission.tenant_id == tenant_id,
+                    )
+                    .first()
+                )
+
+                if not permission:
+                    # no permission set, allow access for everyone
+                    return view(*args, **kwargs)
+
+                if install_required:
+                    if permission.install_permission == TenantPluginPermission.InstallPermission.NOBODY:
+                        raise Forbidden()
+                    if permission.install_permission == TenantPluginPermission.InstallPermission.ADMINS:
+                        if not user.is_admin_or_owner:
+                            raise Forbidden()
+                    if permission.install_permission == TenantPluginPermission.InstallPermission.EVERYONE:
+                        pass
+
+                if debug_required:
+                    if permission.debug_permission == TenantPluginPermission.DebugPermission.NOBODY:
+                        raise Forbidden()
+                    if permission.debug_permission == TenantPluginPermission.DebugPermission.ADMINS:
+                        if not user.is_admin_or_owner:
+                            raise Forbidden()
+                    if permission.debug_permission == TenantPluginPermission.DebugPermission.EVERYONE:
+                        pass
+
+            return view(*args, **kwargs)
+
+        return decorated
+
+    return interceptor