Skip to content

Commit

Permalink
Fix calculator ACE vulnerability
Browse files Browse the repository at this point in the history
  • Loading branch information
Bobholamovic committed Jan 18, 2024
1 parent 0109e37 commit e9cbff3
Showing 1 changed file with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion erniebot-agent/src/erniebot_agent/tools/calculator_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, "<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
return {"formula_result": result}

@property
def examples(self) -> List[Message]:
Expand Down

0 comments on commit e9cbff3

Please sign in to comment.