From e9cbff37febc2b366ed3f4e8e40b8ab7639a54c7 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Thu, 18 Jan 2024 17:25:13 +0800 Subject: [PATCH 1/6] Fix calculator ACE vulnerability --- .../src/erniebot_agent/tools/calculator_tool.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py index e6299b90b..9145a8a74 100644 --- a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py +++ b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py @@ -24,7 +24,16 @@ class CalculatorTool(Tool): ouptut_type: Type[ToolParameterView] = CalculatorToolOutputView async def __call__(self, math_formula: str) -> Dict[str, float]: - return {"formula_result": eval(math_formula)} + try: + code = compile(math_formula, "", "eval") + except (SyntaxError, ValueError) as e: + raise ValueError("Invalid input expression") from e + try: + result = eval(code, {"__builtins__": {}}, {}) + except NameError as e: + names_not_allowed = code.co_names + raise ValueError(f"Names {names_not_allowed} are not allowed in the expression.") from e + return {"formula_result": result} @property def examples(self) -> List[Message]: From 6d07755bce690440ae88d7bd8512aa031c987b89 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Thu, 18 Jan 2024 18:04:42 +0800 Subject: [PATCH 2/6] Add result type check --- erniebot-agent/src/erniebot_agent/tools/calculator_tool.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py index 9145a8a74..7e5ee6702 100644 --- a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py +++ b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py @@ -24,6 +24,9 @@ class CalculatorTool(Tool): ouptut_type: Type[ToolParameterView] = CalculatorToolOutputView async def __call__(self, math_formula: str) -> Dict[str, float]: + # In this eval-based implementation, non-mathematical Python expressions + # are not rejected. Should we do regex checks to ensure that the input + # is a mathematical expression? try: code = compile(math_formula, "", "eval") except (SyntaxError, ValueError) as e: @@ -33,6 +36,8 @@ async def __call__(self, math_formula: str) -> Dict[str, float]: except NameError as e: names_not_allowed = code.co_names raise ValueError(f"Names {names_not_allowed} are not allowed in the expression.") from e + if not isinstance(result, (float, int)): + raise ValueError("The evaluation result of the expression is not a number.") return {"formula_result": result} @property From e1d91376ab5258b796ee89ed4c434f284cc1cd2a Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Thu, 18 Jan 2024 18:05:54 +0800 Subject: [PATCH 3/6] Fix comments --- erniebot-agent/src/erniebot_agent/tools/calculator_tool.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py index 7e5ee6702..57e9268a6 100644 --- a/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py +++ b/erniebot-agent/src/erniebot_agent/tools/calculator_tool.py @@ -24,9 +24,9 @@ class CalculatorTool(Tool): ouptut_type: Type[ToolParameterView] = CalculatorToolOutputView async def __call__(self, math_formula: str) -> Dict[str, float]: - # In this eval-based implementation, non-mathematical Python expressions - # are not rejected. Should we do regex checks to ensure that the input - # is a mathematical expression? + # XXX: In this eval-based implementation, non-mathematical Python + # expressions are not rejected. Should we do regex checks to ensure that + # the input is a mathematical expression? try: code = compile(math_formula, "", "eval") except (SyntaxError, ValueError) as e: From ff8befc2138cb12eca578097b3d45b44371a3e00 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Fri, 19 Jan 2024 18:10:29 +0800 Subject: [PATCH 4/6] Add type ignore error code --- .../applications/erniebot_researcher/tools/utils.py | 2 +- .../src/erniebot_agent/agents/mixins/gradio_mixin.py | 2 +- .../src/erniebot_agent/file/global_file_manager_handler.py | 4 ++-- erniebot-agent/src/erniebot_agent/utils/misc.py | 2 +- erniebot/src/erniebot/utils/bos.py | 6 +++--- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/erniebot-agent/applications/erniebot_researcher/tools/utils.py b/erniebot-agent/applications/erniebot_researcher/tools/utils.py index 5c7bf166f..b88a3ee6e 100644 --- a/erniebot-agent/applications/erniebot_researcher/tools/utils.py +++ b/erniebot-agent/applications/erniebot_researcher/tools/utils.py @@ -6,7 +6,7 @@ from typing import Any, Dict, List, Union import jsonlines -import markdown # type: ignore +import markdown # type: ignore[import-untyped] from langchain.docstore.document import Document from langchain.document_loaders import PyPDFDirectoryLoader from langchain.output_parsers.json import parse_json_markdown diff --git a/erniebot-agent/src/erniebot_agent/agents/mixins/gradio_mixin.py b/erniebot-agent/src/erniebot_agent/agents/mixins/gradio_mixin.py index e7bc23702..08293c865 100644 --- a/erniebot-agent/src/erniebot_agent/agents/mixins/gradio_mixin.py +++ b/erniebot-agent/src/erniebot_agent/agents/mixins/gradio_mixin.py @@ -30,7 +30,7 @@ def launch_gradio_demo(self: BaseAgent, **launch_kwargs: Any): # be constructed outside an event loop, which is probably not sensible. # TODO: Unified optional dependencies management try: - import gradio as gr # type: ignore + import gradio as gr # type: ignore[import-untyped] except ImportError: raise ImportError( "Could not import gradio, which is required for `launch_gradio_demo()`." diff --git a/erniebot-agent/src/erniebot_agent/file/global_file_manager_handler.py b/erniebot-agent/src/erniebot_agent/file/global_file_manager_handler.py index 1d9842eb7..7fa0b09b5 100644 --- a/erniebot-agent/src/erniebot_agent/file/global_file_manager_handler.py +++ b/erniebot-agent/src/erniebot_agent/file/global_file_manager_handler.py @@ -16,14 +16,14 @@ import weakref from typing import Any, NoReturn, Optional, final -import asyncio_atexit # type: ignore +import asyncio_atexit # type: ignore[import-untyped] from typing_extensions import Self from erniebot_agent.file.file_manager import FileManager from erniebot_agent.file.remote_file import AIStudioFileClient from erniebot_agent.utils import config_from_environ as C -_registry = weakref.WeakKeyDictionary() # type: ignore +_registry: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() @final diff --git a/erniebot-agent/src/erniebot_agent/utils/misc.py b/erniebot-agent/src/erniebot_agent/utils/misc.py index 98ca90d20..d145c4ecf 100644 --- a/erniebot-agent/src/erniebot_agent/utils/misc.py +++ b/erniebot-agent/src/erniebot_agent/utils/misc.py @@ -14,7 +14,7 @@ class SingletonMeta(type): - _insts = {} # type: ignore + _insts: dict = {} def __call__(cls, *args, **kwargs): # XXX: We note that the instance created in this way can be actually diff --git a/erniebot/src/erniebot/utils/bos.py b/erniebot/src/erniebot/utils/bos.py index b6cf252df..45dcf9d3a 100644 --- a/erniebot/src/erniebot/utils/bos.py +++ b/erniebot/src/erniebot/utils/bos.py @@ -26,9 +26,9 @@ def upload_file_to_bos( access_key_id: Optional[str] = None, secret_access_key: Optional[str] = None, ) -> str: - from baidubce.auth.bce_credentials import BceCredentials # type: ignore - from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore - from baidubce.services.bos.bos_client import BosClient # type: ignore + from baidubce.auth.bce_credentials import BceCredentials # type: ignore[import-untyped] + from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore[import-untyped] + from baidubce.services.bos.bos_client import BosClient # type: ignore[import-untyped] b_config = BceClientConfiguration( credentials=BceCredentials(access_key_id, secret_access_key), endpoint=bos_host From 027ea87fb9865b7ac5dba7ee8c51f77d6a90f2e1 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Fri, 19 Jan 2024 18:13:23 +0800 Subject: [PATCH 5/6] Fix flake8 and mypy conflicts --- erniebot/src/erniebot/utils/bos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/erniebot/src/erniebot/utils/bos.py b/erniebot/src/erniebot/utils/bos.py index 45dcf9d3a..b6cf252df 100644 --- a/erniebot/src/erniebot/utils/bos.py +++ b/erniebot/src/erniebot/utils/bos.py @@ -26,9 +26,9 @@ def upload_file_to_bos( access_key_id: Optional[str] = None, secret_access_key: Optional[str] = None, ) -> str: - from baidubce.auth.bce_credentials import BceCredentials # type: ignore[import-untyped] - from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore[import-untyped] - from baidubce.services.bos.bos_client import BosClient # type: ignore[import-untyped] + from baidubce.auth.bce_credentials import BceCredentials # type: ignore + from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore + from baidubce.services.bos.bos_client import BosClient # type: ignore b_config = BceClientConfiguration( credentials=BceCredentials(access_key_id, secret_access_key), endpoint=bos_host From 0e898a7c6c6fcd44dfae628e0c38dc63ee19f720 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Fri, 19 Jan 2024 18:15:29 +0800 Subject: [PATCH 6/6] Fix flake8 and mypy conflicts --- erniebot/src/erniebot/utils/bos.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/erniebot/src/erniebot/utils/bos.py b/erniebot/src/erniebot/utils/bos.py index b6cf252df..debf72f20 100644 --- a/erniebot/src/erniebot/utils/bos.py +++ b/erniebot/src/erniebot/utils/bos.py @@ -26,9 +26,15 @@ def upload_file_to_bos( access_key_id: Optional[str] = None, secret_access_key: Optional[str] = None, ) -> str: - from baidubce.auth.bce_credentials import BceCredentials # type: ignore - from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore - from baidubce.services.bos.bos_client import BosClient # type: ignore + from baidubce.auth.bce_credentials import ( # type: ignore[import-untyped] + BceCredentials, + ) + from baidubce.bce_client_configuration import ( # type: ignore[import-untyped] + BceClientConfiguration, + ) + from baidubce.services.bos.bos_client import ( # type: ignore[import-untyped] + BosClient, + ) b_config = BceClientConfiguration( credentials=BceCredentials(access_key_id, secret_access_key), endpoint=bos_host