mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-22 03:41:43 +00:00
fixed liting errors
This commit is contained in:
parent
9c219db9ba
commit
dc8bdae64f
@ -11,7 +11,6 @@ from concurrent.futures import Executor, ThreadPoolExecutor
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, final
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, final
|
||||||
|
|
||||||
from jinja2 import Template
|
|
||||||
from jinja2.sandbox import SandboxedEnvironment
|
from jinja2.sandbox import SandboxedEnvironment
|
||||||
|
|
||||||
from dbgpt._private.pydantic import ConfigDict, Field
|
from dbgpt._private.pydantic import ConfigDict, Field
|
||||||
@ -40,18 +39,18 @@ class ConversableAgent(Role, Agent):
|
|||||||
"""ConversableAgent is an agent that can communicate with other agents."""
|
"""ConversableAgent is an agent that can communicate with other agents."""
|
||||||
|
|
||||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||||
|
|
||||||
# Dangerous template patterns that could lead to code execution
|
# Dangerous template patterns that could lead to code execution
|
||||||
_DANGEROUS_TEMPLATE_PATTERNS = [
|
_DANGEROUS_TEMPLATE_PATTERNS = [
|
||||||
r'\{\{.*__.*\}\}', # Double underscore methods
|
r"\{\{.*__.*\}\}", # Double underscore methods
|
||||||
r'\{\{.*import.*\}\}', # Import statements
|
r"\{\{.*import.*\}\}", # Import statements
|
||||||
r'\{\{.*exec.*\}\}', # Exec calls
|
r"\{\{.*exec.*\}\}", # Exec calls
|
||||||
r'\{\{.*eval.*\}\}', # Eval calls
|
r"\{\{.*eval.*\}\}", # Eval calls
|
||||||
r'\{\{.*open.*\}\}', # File operations
|
r"\{\{.*open.*\}\}", # File operations
|
||||||
r'\{\{.*subprocess.*\}\}', # Subprocess calls
|
r"\{\{.*subprocess.*\}\}", # Subprocess calls
|
||||||
r'\{\{.*os\..*\}\}', # OS module access
|
r"\{\{.*os\..*\}\}", # OS module access
|
||||||
r'\{\{.*globals.*\}\}', # Globals access
|
r"\{\{.*globals.*\}\}", # Globals access
|
||||||
r'\{\{.*\[.*\].*\}\}', # Bracket notation access
|
r"\{\{.*\[.*\].*\}\}", # Bracket notation access
|
||||||
]
|
]
|
||||||
|
|
||||||
agent_context: Optional[AgentContext] = Field(None, description="Agent context")
|
agent_context: Optional[AgentContext] = Field(None, description="Agent context")
|
||||||
@ -902,13 +901,13 @@ class ConversableAgent(Role, Agent):
|
|||||||
"""Sanitize template parameters to prevent injection attacks."""
|
"""Sanitize template parameters to prevent injection attacks."""
|
||||||
if not params:
|
if not params:
|
||||||
return params
|
return params
|
||||||
|
|
||||||
return self._sanitize_dict(params)
|
return self._sanitize_dict(params)
|
||||||
|
|
||||||
def _sanitize_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
def _sanitize_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||||
"""Sanitize dictionary values recursively."""
|
"""Sanitize dictionary values recursively."""
|
||||||
return {key: self._sanitize_value(value) for key, value in data.items()}
|
return {key: self._sanitize_value(value) for key, value in data.items()}
|
||||||
|
|
||||||
def _sanitize_value(self, value: Any) -> Any:
|
def _sanitize_value(self, value: Any) -> Any:
|
||||||
"""Sanitize a single value based on its type."""
|
"""Sanitize a single value based on its type."""
|
||||||
if isinstance(value, str):
|
if isinstance(value, str):
|
||||||
@ -918,7 +917,7 @@ class ConversableAgent(Role, Agent):
|
|||||||
elif isinstance(value, list):
|
elif isinstance(value, list):
|
||||||
return [self._sanitize_value(item) for item in value]
|
return [self._sanitize_value(item) for item in value]
|
||||||
return value
|
return value
|
||||||
|
|
||||||
def _sanitize_string(self, text: str) -> str:
|
def _sanitize_string(self, text: str) -> str:
|
||||||
"""Check string for dangerous template injection patterns."""
|
"""Check string for dangerous template injection patterns."""
|
||||||
for pattern in self._DANGEROUS_TEMPLATE_PATTERNS:
|
for pattern in self._DANGEROUS_TEMPLATE_PATTERNS:
|
||||||
@ -1100,7 +1099,7 @@ class ConversableAgent(Role, Agent):
|
|||||||
"""Build system prompt with security controls."""
|
"""Build system prompt with security controls."""
|
||||||
if self.bind_prompt:
|
if self.bind_prompt:
|
||||||
return self._render_bind_prompt(resource_vars, context)
|
return self._render_bind_prompt(resource_vars, context)
|
||||||
|
|
||||||
# Fallback to build_prompt with sanitized context
|
# Fallback to build_prompt with sanitized context
|
||||||
sanitized_context = self.sanitize_template_params(context or {})
|
sanitized_context = self.sanitize_template_params(context or {})
|
||||||
return await self.build_prompt(
|
return await self.build_prompt(
|
||||||
@ -1111,11 +1110,11 @@ class ConversableAgent(Role, Agent):
|
|||||||
is_retry_chat=is_retry_chat,
|
is_retry_chat=is_retry_chat,
|
||||||
**sanitized_context,
|
**sanitized_context,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _render_bind_prompt(
|
def _render_bind_prompt(
|
||||||
self,
|
self,
|
||||||
resource_vars: Optional[Dict] = None,
|
resource_vars: Optional[Dict] = None,
|
||||||
context: Optional[Dict[str, Any]] = None
|
context: Optional[Dict[str, Any]] = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Render bind prompt template with sanitized parameters."""
|
"""Render bind prompt template with sanitized parameters."""
|
||||||
prompt_param = {}
|
prompt_param = {}
|
||||||
@ -1124,7 +1123,7 @@ class ConversableAgent(Role, Agent):
|
|||||||
if context:
|
if context:
|
||||||
sanitized_context = self.sanitize_template_params(context)
|
sanitized_context = self.sanitize_template_params(context)
|
||||||
prompt_param.update(sanitized_context)
|
prompt_param.update(sanitized_context)
|
||||||
|
|
||||||
if self.bind_prompt.template_format == "f-string":
|
if self.bind_prompt.template_format == "f-string":
|
||||||
return self.bind_prompt.template.format(**prompt_param)
|
return self.bind_prompt.template.format(**prompt_param)
|
||||||
elif self.bind_prompt.template_format == "jinja2":
|
elif self.bind_prompt.template_format == "jinja2":
|
||||||
|
@ -31,7 +31,7 @@ def _jinja2_formatter(template: str, **kwargs: Any) -> str:
|
|||||||
"jinja2 not installed, which is needed to use the jinja2_formatter. "
|
"jinja2 not installed, which is needed to use the jinja2_formatter. "
|
||||||
"Please install it with `pip install jinja2`."
|
"Please install it with `pip install jinja2`."
|
||||||
)
|
)
|
||||||
|
|
||||||
env = SandboxedEnvironment()
|
env = SandboxedEnvironment()
|
||||||
return env.from_string(template).render(**kwargs)
|
return env.from_string(template).render(**kwargs)
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user