sangge/BackDoorBuster
sangge/BackDoorBuster
1
0
Fork 0
Code Issues Pull Requests Projects 1 Activity

feature/GPT #12

Merged
sangge merged 18 commits from feature/GPT into main 2024-04-29 18:58:49 +08:00
Conversation 3 Commits 18 Files Changed 9 +14 -11
1 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 4abd93f688 - Show all commits

23
detection/GPTdetection.py
View File

@ -1,7 +1,5 @@
import json import json
import os import os
import re
import sys
from .utils import * from .utils import *
import openai import openai
@ -15,22 +13,27 @@ def detect_gpt(filename: str):
{ {
"role": "system", "role": "system",
"content": "You are a Python code reviewer.Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: " "content": "You are a Python code reviewer.Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: "
"[{\"Line\": {the line number}, \"Risk\": \"{choose from (high,medium,low)}\",\"Reason\":\"{how it is vulnable}\"}] Each of these three field is required.\n" '[{"Line": {the line number}, "Risk": "{choose from (high,medium,low)}","Reason":"{how it is vulnable}"}] Each of these three field is required.\n'
"You are required to only output the json format. Do not output any other information.\n" "You are required to only output the json format. Do not output any other information.\n",
}, },
{ {
"role": "user", "role": "user",
"content": text, "content": text,
} },
], ],
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
) )
try: try:
res_json = json.loads(response.choices[0].message.content) message_content = response.choices[0].message.content
except: if message_content is None:
print("Error: Could not parse the response. Please try again.") raise ValueError("API response content is None")
sys.exit(1) res_json = json.loads(message_content)
except json.JSONDecodeError:
raise ValueError("Error: Could not parse the response. Please try again.")
classified_results = {"high": [], "medium": [], "low": [], "none": []} classified_results = {"high": [], "medium": [], "low": [], "none": []}
for res in res_json: for res in res_json:
classified_results[res["Risk"]].append((res["Line"], text.split("\n")[res["Line"] - 1].strip())) classified_results[res["Risk"]].append(
(res["Line"], text.split("\n")[res["Line"] - 1].strip())
)
return classified_results return classified_results