fix:修复文心一言的调用

This commit is contained in:
2024-06-05 10:36:26 +08:00
parent e9b1e82492
commit 167bbe0a14
4 changed files with 62 additions and 56 deletions

View File

@@ -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]
]
)