diff --git a/astrbot/core/computer/tools/fs.py b/astrbot/core/computer/tools/fs.py index 126da42586..31b7f3f513 100644 --- a/astrbot/core/computer/tools/fs.py +++ b/astrbot/core/computer/tools/fs.py @@ -11,6 +11,7 @@ from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from ..computer_client import get_booter +from .permissions import check_admin_permission # @dataclass # class CreateFileTool(FunctionTool): @@ -102,6 +103,8 @@ async def call( context: ContextWrapper[AstrAgentContext], local_path: str, ) -> str | None: + if permission_error := check_admin_permission(context, "File upload/download"): + return permission_error sb = await get_booter( context.context.context, context.context.event.unified_msg_origin, @@ -161,6 +164,8 @@ async def call( remote_path: str, also_send_to_user: bool = True, ) -> ToolExecResult: + if permission_error := check_admin_permission(context, "File upload/download"): + return permission_error sb = await get_booter( context.context.context, context.context.event.unified_msg_origin, diff --git a/astrbot/core/computer/tools/permissions.py b/astrbot/core/computer/tools/permissions.py new file mode 100644 index 0000000000..489f485f9d --- /dev/null +++ b/astrbot/core/computer/tools/permissions.py @@ -0,0 +1,19 @@ +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.astr_agent_context import AstrAgentContext + + +def check_admin_permission( + context: ContextWrapper[AstrAgentContext], operation_name: str +) -> str | None: + cfg = context.context.context.get_config( + umo=context.context.event.unified_msg_origin + ) + provider_settings = cfg.get("provider_settings", {}) + require_admin = provider_settings.get("computer_use_require_admin", True) + if require_admin and context.context.event.role != "admin": + return ( + f"error: Permission denied. {operation_name} is only allowed for admin users. " + "Tell user to set admins in `AstrBot WebUI -> Config -> General Config` by adding their user ID to the admins list if they need this feature. " + f"User's ID is: {context.context.event.get_sender_id()}. User's ID can be found by using /sid command." + ) + return None diff --git a/astrbot/core/computer/tools/python.py b/astrbot/core/computer/tools/python.py index 26eaba8e3a..cc835bc753 100644 --- a/astrbot/core/computer/tools/python.py +++ b/astrbot/core/computer/tools/python.py @@ -7,6 +7,7 @@ from astrbot.core.agent.tool import ToolExecResult from astrbot.core.astr_agent_context import AstrAgentContext, AstrMessageEvent from astrbot.core.computer.computer_client import get_booter, get_local_booter +from astrbot.core.computer.tools.permissions import check_admin_permission from astrbot.core.message.message_event_result import MessageChain param_schema = { @@ -26,21 +27,6 @@ } -def _check_admin_permission(context: ContextWrapper[AstrAgentContext]) -> str | None: - cfg = context.context.context.get_config( - umo=context.context.event.unified_msg_origin - ) - provider_settings = cfg.get("provider_settings", {}) - require_admin = provider_settings.get("computer_use_require_admin", True) - if require_admin and context.context.event.role != "admin": - return ( - "error: Permission denied. Python execution is only allowed for admin users. " - "Tell user to set admins in `AstrBot WebUI -> Config -> General Config` by adding their user ID to the admins list if they need this feature." - f"User's ID is: {context.context.event.get_sender_id()}. User's ID can be found by using /sid command." - ) - return None - - async def handle_result(result: dict, event: AstrMessageEvent) -> ToolExecResult: data = result.get("data", {}) output = data.get("output", {}) @@ -81,7 +67,7 @@ class PythonTool(FunctionTool): async def call( self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False ) -> ToolExecResult: - if permission_error := _check_admin_permission(context): + if permission_error := check_admin_permission(context, "Python execution"): return permission_error sb = await get_booter( context.context.context, @@ -104,7 +90,7 @@ class LocalPythonTool(FunctionTool): async def call( self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False ) -> ToolExecResult: - if permission_error := _check_admin_permission(context): + if permission_error := check_admin_permission(context, "Python execution"): return permission_error sb = get_local_booter() try: diff --git a/astrbot/core/computer/tools/shell.py b/astrbot/core/computer/tools/shell.py index b07f98a212..9e729573a7 100644 --- a/astrbot/core/computer/tools/shell.py +++ b/astrbot/core/computer/tools/shell.py @@ -7,21 +7,7 @@ from astrbot.core.astr_agent_context import AstrAgentContext from ..computer_client import get_booter, get_local_booter - - -def _check_admin_permission(context: ContextWrapper[AstrAgentContext]) -> str | None: - cfg = context.context.context.get_config( - umo=context.context.event.unified_msg_origin - ) - provider_settings = cfg.get("provider_settings", {}) - require_admin = provider_settings.get("computer_use_require_admin", True) - if require_admin and context.context.event.role != "admin": - return ( - "error: Permission denied. Shell execution is only allowed for admin users. " - "Tell user to set admins in `AstrBot WebUI -> Config -> General Config` by adding their user ID to the admins list if they need this feature." - f"User's ID is: {context.context.event.get_sender_id()}. User's ID can be found by using /sid command." - ) - return None +from .permissions import check_admin_permission @dataclass @@ -61,7 +47,7 @@ async def call( background: bool = False, env: dict = {}, ) -> ToolExecResult: - if permission_error := _check_admin_permission(context): + if permission_error := check_admin_permission(context, "Shell execution"): return permission_error if self.is_local: