feat:对正则匹配出的文件再过一次llm检测

This commit is contained in:
tritium0041 2024-06-03 20:34:30 +08:00
parent e653ddd726
commit 4e67f4ebed
2 changed files with 20 additions and 2 deletions

View File

@ -25,7 +25,7 @@ def detectGPT(content: str):
signal.signal(signal.SIGTERM, timeout_handler) signal.signal(signal.SIGTERM, timeout_handler)
signal.alarm(10) signal.alarm(10)
client = openai.OpenAI(api_key=api_key) client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1",api_key=api_key)
text = content text = content
# client = openai.OpenAI(api_key="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") #测试用key # client = openai.OpenAI(api_key="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") #测试用key
response = client.chat.completions.create( response = client.chat.completions.create(

View File

@ -11,23 +11,41 @@ from detection.GPTdetection import detectGPT
class TestFinalTests(unittest.TestCase): class TestFinalTests(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
shutil.rmtree("./tmp/repo", ignore_errors=True) shutil.rmtree("./tmp/repo", ignore_errors=True)
clone_repo("https://github.com/TheAlgorithms/Python.git", "./tmp/repo") clone_repo("https://github.com/injetlee/Python.git", "./tmp/repo")
sampleRate = 0.1 sampleRate = 0.1
self.injectedNum = inject_random_backdoor("./tmp/repo", sample_rate=sampleRate) self.injectedNum = inject_random_backdoor("./tmp/repo", sample_rate=sampleRate)
print(self.injectedNum)
project_path = Path("./tmp/repo") project_path = Path("./tmp/repo")
self.all_python_files = list(project_path.rglob("*.py")) self.all_python_files = list(project_path.rglob("*.py"))
self.filesNum = len(self.all_python_files) self.filesNum = len(self.all_python_files)
self.trueRate = self.injectedNum / self.filesNum self.trueRate = self.injectedNum / self.filesNum
print(self.trueRate)
def test_final_tests(self): def test_final_tests(self):
detectedNum = 0 detectedNum = 0
possibly_dangerous_file = []
for file in self.all_python_files: for file in self.all_python_files:
content = read_file_content(str(file)) content = read_file_content(str(file))
results = find_dangerous_functions(content, ".py") results = find_dangerous_functions(content, ".py")
if len(results["high"]) > 0 or len(results["medium"]) > 0 or len(results["low"]) > 0: if len(results["high"]) > 0 or len(results["medium"]) > 0 or len(results["low"]) > 0:
detectedNum += 1 detectedNum += 1
possibly_dangerous_file.append(file)
print(detectedNum / self.filesNum) print(detectedNum / self.filesNum)
self.assertAlmostEqual(detectedNum / self.filesNum, self.trueRate, places=1) self.assertAlmostEqual(detectedNum / self.filesNum, self.trueRate, places=1)
GPTdetectedNum = 0
for i in possibly_dangerous_file:
content = read_file_content(str(i))
results = {}
try:
results = detectGPT(content)
except Exception as e:
print(e)
if len(results["high"]) > 0 or len(results["medium"]) > 0 or len(results["low"]) > 0:
GPTdetectedNum += 1
print(GPTdetectedNum)
self.assertGreaterEqual(GPTdetectedNum, detectedNum)