Fix CalculatorTool ACE vulnerability (#301)

* Fix calculator ACE vulnerability

* Add result type check

* Fix comments

* Add type ignore error code

* Fix flake8 and mypy conflicts

* Fix flake8 and mypy conflicts

Author: Bobholamovic

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.output_parsers.json import parse_json_markdown
 from weasyprint import CSS, HTML

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()`."

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

erniebot-agent/src/erniebot_agent/tools/calculator_tool.py

@@ -24,7 +24,21 @@ class CalculatorTool(Tool):
     output_type: Type[ToolParameterView] = CalculatorToolOutputView
 
     async def __call__(self, math_formula: str) -> Dict[str, float]:
-        return {"formula_result": eval(math_formula)}
+        # 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, "<string>", "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
+        if not isinstance(result, (float, int)):
+            raise ValueError("The evaluation result of the expression is not a number.")
+        return {"formula_result": result}
 
     @property
     def examples(self) -> List[Message]:

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

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