fix：修复文心一言的调用

detection/GPTdetection.py

@@ -92,6 +92,7 @@ def GPTdetectFileList(fileList):
 def GPTThread(filename, content, results):
     try:
         res = detectGPT(content)
+        # print(res)
         for key in res:
             if key != "none":  # Exclude 'none' risk level
                 results[key].extend(

detection/__main__.py

@@ -8,6 +8,7 @@ from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
 from detection.pickle_detection import pickleDataDetection
 from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT,GPTdetectFileList
+# from .cngptdetection import detectGPT,GPTdetectFileList
 from .pyc_detection import disassemble_pyc
 from .utils import *
 import sys

detection/cngptdetection.py

@@ -1,16 +1,21 @@
 import os
+import threading
+import time
+
 import requests
 import re
 import json
 from typing import List, Dict, Any
 
+from detection.utils import read_file_content
+
 
 class TimeoutException(Exception):
     """自定义异常用于处理超时情况。"""
     pass
 
 
-def detectGPT(content: str) -> str:
+def detectGPT(content: str,token:str):
     """
     检测给定的代码内容中的潜在安全漏洞。
 
@@ -20,15 +25,8 @@ def detectGPT(content: str) -> str:
     返回:
     - 分类后的漏洞信息的JSON字符串。
     """
-    api_key = os.getenv("BAIDU_API_KEY")
-    secret_key = os.getenv("BAIDU_SECRET_KEY")
-    #api_key = "DUBWNIrB6QJLOsLkpnEz2ZZa"
-    #secret_key = "9WK4HIV2n9r1ePPirqD4EQ6Ea33rH1m7"
-    if not api_key or not secret_key:
-        raise ValueError("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
 
-    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-8k-0329?access_token=" + get_access_token(
-        api_key, secret_key)
+    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-4.0-8k-0329?access_token=" + token
 
     payload = json.dumps({
         "messages": [
@@ -63,6 +61,7 @@ def detectGPT(content: str) -> str:
 
     classified_results = {"high": [], "medium": [], "low": [], "none": []}
     for res in extracted_data:
+        # print(res)
         try:
             line_number = int(res["Line"])
             classified_results[res["Risk"]].append(
@@ -71,7 +70,7 @@ def detectGPT(content: str) -> str:
         except (ValueError, IndexError, KeyError):
             continue
 
-    return json.dumps(classified_results, indent=2, ensure_ascii=False)
+    return classified_results
 
 
 def get_access_token(api_key: str, secret_key: str) -> str:
@@ -110,4 +109,41 @@ def extract_json_from_text(text: str) -> List[Dict[str, Any]]:
         print(f"解码 JSON 时出错: {e}")
         return []
 
-    return data
+    return data
+
+
+def GPTdetectFileList(fileList):
+    api_key = os.getenv("BAIDU_API_KEY")
+    secret_key = os.getenv("BAIDU_SECRET_KEY")
+    # api_key = "DUBWNIrB6QJLOsLkpnEz2ZZa"
+    # secret_key = "9WK4HIV2n9r1ePPirqD4EQ6Ea33rH1m7"
+    if not api_key or not secret_key:
+        raise ValueError("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
+    # print(len(fileList))
+    results = {"high": [], "medium": [], "low": [], "none": []}
+    threads = []
+    token = get_access_token(api_key, secret_key)
+    # print(token)
+    for file in fileList:
+        content = read_file_content(str(file))
+        threads.append(threading.Thread(target=GPTThread, args=(str(file), content, results,token)))
+    for thread in threads:
+        thread.start()
+        time.sleep(0.5)
+    for thread in threads:
+        thread.join()
+    return results
+
+
+def GPTThread(filename, content, results,token):
+
+        res = detectGPT(content,token)
+        # print(res)
+        for key in res:
+            if key != "none":  # Exclude 'none' risk level
+                results[key].extend(
+                    [
+                        (f"{filename}: Line {line_num}", line)
+                        for line_num, line in res[key]
+                    ]
+                )