From 0e0ab7cab2d85b52509d99075cbbd27058203d38 Mon Sep 17 00:00:00 2001 From: nkoorty Date: Mon, 7 Jul 2025 15:26:17 +0200 Subject: [PATCH 1/2] fix: remediated idor --- .../src/dbgpt_serve/agent/hub/controller.py | 98 +++++++++++++++++-- 1 file changed, 88 insertions(+), 10 deletions(-) diff --git a/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py b/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py index 6763543c9..5cedc4455 100644 --- a/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py +++ b/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py @@ -1,8 +1,10 @@ import logging from abc import ABC -from typing import List +from functools import cache +from typing import List, Optional -from fastapi import APIRouter, Body, File, UploadFile +from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile +from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer from dbgpt.agent.resource.tool.autogpt.plugins_util import scan_plugins from dbgpt.agent.resource.tool.pack import AutoGPTPluginToolPack @@ -24,6 +26,55 @@ from .model.model import MyPluginVO, PluginHubVO router = APIRouter() logger = logging.getLogger(__name__) +get_bearer_token = HTTPBearer(auto_error=False) +_api_keys_config: Optional[str] = None + + +def set_api_keys_config(api_keys: Optional[str]): + global _api_keys_config + _api_keys_config = api_keys + + +@cache +def _parse_api_keys(api_keys: str) -> List[str]: + if not api_keys: + return [] + return [key.strip() for key in api_keys.split(",")] + + +async def check_api_key( + auth: Optional[HTTPAuthorizationCredentials] = Depends(get_bearer_token), +) -> Optional[str]: + global _api_keys_config + if _api_keys_config: + api_keys = _parse_api_keys(_api_keys_config) + if auth is None or (token := auth.credentials) not in api_keys: + raise HTTPException( + status_code=401, + detail={ + "error": { + "message": "", + "type": "invalid_request_error", + "param": None, + "code": "invalid_api_key", + } + }, + ) + return token + return None + + +def _validate_user_access(requested_user: Optional[str], authenticated_user: Optional[str]) -> str: + global _api_keys_config + if _api_keys_config and authenticated_user: + if requested_user and requested_user != authenticated_user: + raise HTTPException( + status_code=403, + detail="Access denied: Cannot access another user's resources" + ) + return requested_user or authenticated_user + return requested_user or "default" + class ModulePlugin(BaseComponent, ABC): name = ComponentType.PLUGIN_HUB @@ -34,6 +85,14 @@ class ModulePlugin(BaseComponent, ABC): def init_app(self, system_app: SystemApp): system_app.app.include_router(router, prefix="/api", tags=["Agent"]) + try: + global_api_keys = system_app.config.get("dbgpt.app.global.api_keys") + if global_api_keys: + set_api_keys_config(global_api_keys) + elif hasattr(system_app, 'config') and hasattr(system_app.config, 'API_KEYS'): + set_api_keys_config(system_app.config.API_KEYS) + except Exception: + pass def refresh_plugins(self): self.plugins = scan_plugins(PLUGINS_DIR) @@ -89,18 +148,27 @@ async def get_agent_list(filter: PagenationFilter[PluginHubFilter] = Body()): @router.post("/v1/agent/my", response_model=Result[List[MyPluginVO]]) -async def my_agents(user: str = None): +async def my_agents( + user: str = None, + authenticated_user: Optional[str] = Depends(check_api_key) +): logger.info(f"my_agents:{user}") - agents = plugin_hub.get_my_plugin(user) + effective_user = _validate_user_access(user, authenticated_user) + agents = plugin_hub.get_my_plugin(effective_user) agent_dicts = MyPluginEntity.to_vo(agents) return Result.succ(agent_dicts) @router.post("/v1/agent/install", response_model=Result[str]) -async def agent_install(plugin_name: str, user: str = None): +async def agent_install( + plugin_name: str, + user: str = None, + authenticated_user: Optional[str] = Depends(check_api_key) +): logger.info(f"agent_install:{plugin_name},{user}") + effective_user = _validate_user_access(user, authenticated_user) try: - plugin_hub.install_plugin(plugin_name, user) + plugin_hub.install_plugin(plugin_name, effective_user) module_plugin.refresh_plugins() @@ -111,10 +179,15 @@ async def agent_install(plugin_name: str, user: str = None): @router.post("/v1/agent/uninstall", response_model=Result[str]) -async def agent_uninstall(plugin_name: str, user: str = None): +async def agent_uninstall( + plugin_name: str, + user: str = None, + authenticated_user: Optional[str] = Depends(check_api_key) +): logger.info(f"agent_uninstall:{plugin_name},{user}") + effective_user = _validate_user_access(user, authenticated_user) try: - plugin_hub.uninstall_plugin(plugin_name, user) + plugin_hub.uninstall_plugin(plugin_name, effective_user) module_plugin.refresh_plugins() return Result.succ(None) @@ -124,10 +197,15 @@ async def agent_uninstall(plugin_name: str, user: str = None): @router.post("/v1/personal/agent/upload", response_model=Result[str]) -async def personal_agent_upload(doc_file: UploadFile = File(...), user: str = None): +async def personal_agent_upload( + doc_file: UploadFile = File(...), + user: str = None, + authenticated_user: Optional[str] = Depends(check_api_key) +): logger.info(f"personal_agent_upload:{doc_file.filename},{user}") + effective_user = _validate_user_access(user, authenticated_user) try: - await plugin_hub.upload_my_plugin(doc_file, user) + await plugin_hub.upload_my_plugin(doc_file, effective_user) module_plugin.refresh_plugins() return Result.succ(None) except Exception as e: From 81d4dfd0d6b21b6a05488805854b6f9f14527aab Mon Sep 17 00:00:00 2001 From: nkoorty Date: Mon, 7 Jul 2025 18:05:46 +0200 Subject: [PATCH 2/2] fixed linting errors --- .../src/dbgpt_serve/agent/hub/controller.py | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py b/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py index 5cedc4455..058851a43 100644 --- a/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py +++ b/packages/dbgpt-serve/src/dbgpt_serve/agent/hub/controller.py @@ -64,13 +64,15 @@ async def check_api_key( return None -def _validate_user_access(requested_user: Optional[str], authenticated_user: Optional[str]) -> str: +def _validate_user_access( + requested_user: Optional[str], authenticated_user: Optional[str] +) -> str: global _api_keys_config if _api_keys_config and authenticated_user: if requested_user and requested_user != authenticated_user: raise HTTPException( - status_code=403, - detail="Access denied: Cannot access another user's resources" + status_code=403, + detail="Access denied: Cannot access another user's resources", ) return requested_user or authenticated_user return requested_user or "default" @@ -89,7 +91,9 @@ class ModulePlugin(BaseComponent, ABC): global_api_keys = system_app.config.get("dbgpt.app.global.api_keys") if global_api_keys: set_api_keys_config(global_api_keys) - elif hasattr(system_app, 'config') and hasattr(system_app.config, 'API_KEYS'): + elif hasattr(system_app, "config") and hasattr( + system_app.config, "API_KEYS" + ): set_api_keys_config(system_app.config.API_KEYS) except Exception: pass @@ -149,8 +153,7 @@ async def get_agent_list(filter: PagenationFilter[PluginHubFilter] = Body()): @router.post("/v1/agent/my", response_model=Result[List[MyPluginVO]]) async def my_agents( - user: str = None, - authenticated_user: Optional[str] = Depends(check_api_key) + user: str = None, authenticated_user: Optional[str] = Depends(check_api_key) ): logger.info(f"my_agents:{user}") effective_user = _validate_user_access(user, authenticated_user) @@ -161,9 +164,9 @@ async def my_agents( @router.post("/v1/agent/install", response_model=Result[str]) async def agent_install( - plugin_name: str, + plugin_name: str, user: str = None, - authenticated_user: Optional[str] = Depends(check_api_key) + authenticated_user: Optional[str] = Depends(check_api_key), ): logger.info(f"agent_install:{plugin_name},{user}") effective_user = _validate_user_access(user, authenticated_user) @@ -180,9 +183,9 @@ async def agent_install( @router.post("/v1/agent/uninstall", response_model=Result[str]) async def agent_uninstall( - plugin_name: str, + plugin_name: str, user: str = None, - authenticated_user: Optional[str] = Depends(check_api_key) + authenticated_user: Optional[str] = Depends(check_api_key), ): logger.info(f"agent_uninstall:{plugin_name},{user}") effective_user = _validate_user_access(user, authenticated_user) @@ -198,9 +201,9 @@ async def agent_uninstall( @router.post("/v1/personal/agent/upload", response_model=Result[str]) async def personal_agent_upload( - doc_file: UploadFile = File(...), + doc_file: UploadFile = File(...), user: str = None, - authenticated_user: Optional[str] = Depends(check_api_key) + authenticated_user: Optional[str] = Depends(check_api_key), ): logger.info(f"personal_agent_upload:{doc_file.filename},{user}") effective_user = _validate_user_access(user, authenticated_user)