From 28ae4d80aaf260619c6c757776642da80173ca24 Mon Sep 17 00:00:00 2001
From: HCha <hcha@hcha.top>
Date: Wed, 9 Jul 2025 23:09:48 +0800
Subject: [PATCH] fix(function_calling): improve security of function calls

---
 main.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index af62807..2e7b2ae 100644
--- a/main.py
+++ b/main.py
@@ -8,6 +8,14 @@ import socks
 import function_generator as fun
 
 SUPPORTED_MODELS = ["gpt-3.5-turbo-0613", "gpt-4-0613"]
+ALLOWED_FUNCTIONS = {
+    "get_time",
+    "run_cmd",
+    "wolframalpha",
+    "spider",
+    "read",
+    "write",
+}
 
 def set_proxy(HTTP_PROXY = None, SOCKS_PROXY = None):
     """
@@ -84,7 +92,15 @@ def chat(messages, api_key, model, functions, base):
             if use_function:
                 if function_call["name"] != "exit":
                     print(f"\n正在调用插件：{function_call['name']}")
-                function_response = eval(f"fun.{function_call['name']}({function_call['arguments']})")
+                
+                function_name = function_call['name']
+                if function_name in ALLOWED_FUNCTIONS:
+                    function_to_call = getattr(fun, function_name)
+                    # Use .get() for safer access to arguments
+                    function_response = function_to_call(function_call.get('arguments'))
+                else:
+                    function_response = json.dumps({"error": f"Function '{function_name}' is not an allowed function."})
+
                 add_msg(messages, "function", function_response, function_call["name"])
         except Exception as e:
             print(f"调用OpenAI API时发生了错误: {str(e)}")