fix: fix some error

This commit is contained in:
sangge-redmi 2024-06-04 19:27:56 +08:00
parent 610e35f868
commit fd4ecce710
2 changed files with 31 additions and 23 deletions

View File

@ -25,7 +25,7 @@ def clone_repo(repo_url: str, clone_dir: str) -> None:
def inject_random_backdoor(
path: str, pickle: bool = False, pyc: bool = False, sample_rate: float = 0.1
path: str, sample_rate: float = 0.1
) -> Tuple[Tuple[str, int], ...]:
"""
Insert random backdoor into the path.
@ -35,11 +35,6 @@ def inject_random_backdoor(
pickle (bool): Whether to insert a backdoor into a pickle file.
pyc (bool): Whether to insert a backdoor into a compiled Python file.
"""
if pickle:
inject_pickle_backdoor(path)
if pyc:
inject_pyc_backdoor(path)
project_path = Path(path)
all_python_files = list(project_path.rglob("*.py"))
injected_python_files = []
@ -187,4 +182,5 @@ if __name__ == "__main__":
repo_url = "https://github.com/TheAlgorithms/Python.git"
clone_dir = "/tmp/repo"
clone_repo(repo_url, clone_dir)
inject_random_backdoor(clone_dir, pickle=True, pyc=True)
inject_random_backdoor(clone_dir)
inject_pickle_backdoor(clone_dir)

View File

@ -5,7 +5,13 @@ import os
import threading
from detection.utils import read_file_content
from .final_tests_util import clone_repo, Path, inject_random_backdoor
from .final_tests_util import (
clone_repo,
Path,
inject_pickle_backdoor,
inject_random_backdoor,
inject_pyc_backdoor,
)
from detection.Regexdetection import find_dangerous_functions
from detection.GPTdetection import detectGPT
@ -23,12 +29,14 @@ def GPTdetectFileList(fileList):
thread.join()
return results
def GPTThread(content, results):
try:
results.append(detectGPT(content))
except Exception as e:
print(e)
class TestFinalTests(unittest.TestCase):
def setUp(self) -> None:
self.path = "./tmp/repo/"
@ -37,21 +45,21 @@ class TestFinalTests(unittest.TestCase):
clone_repo("https://github.com/TheAlgorithms/Python.git", "/tmp/Python")
shutil.copytree("/tmp/Python", self.path)
sampleRate = 0.1
self.inject_result = inject_random_backdoor(
self.path, sample_rate=sampleRate, pyc=True, pickle=True
)
self.inject_result = inject_random_backdoor(self.path, sample_rate=sampleRate)
self.pickle_true_num = inject_pickle_backdoor(self.path)
self.pyc_true_num = inject_pyc_backdoor(self.path)
self.injectedNum = len(self.inject_result)
print(self.injectedNum)
project_path = Path(self.path)
self.all_python_files = list(project_path.rglob("*.py"))
self.py_filesNum = len(self.all_python_files)
self.py_files_num = len(self.all_python_files)
all_pickle_files = list(project_path.rglob("*.pickle"))
self.pickle_filesNum = len(all_pickle_files)
self.pickle_files_num = len(all_pickle_files)
all_pyc_files = list(project_path.rglob("*.pyc"))
self.pyc_filesNum = len(all_pyc_files)
self.pyc_files_num = len(all_pyc_files)
os.system(
"python -m detection " + self.path + " -o " + self.path + "output.txt"
@ -71,7 +79,7 @@ class TestFinalTests(unittest.TestCase):
):
detectedNum += 1
possibly_dangerous_file.append(file)
print(detectedNum / self.py_filesNum)
print(detectedNum / self.py_files_num)
GPTdetectedNum = 0
for i in possibly_dangerous_file:
@ -100,12 +108,16 @@ class TestFinalTests(unittest.TestCase):
injected_detectedNum += 1
injected_accurency = injected_detectedNum / self.injectedNum
print(f"injected files accurency: {injected_accurency}")
try:
GPTresult = GPTdetectFileList(possibly_dangerous_file)
for result in GPTresult:
if len(result) > 0:
GPTdetectedNum += 1
print(GPTdetectedNum)
self.assertGreaterEqual(GPTdetectedNum, detectedNum)
except Exception as e:
# print(e)
pass
# test pickle files
with open(self.path + "output.txt", "r") as f:
@ -114,7 +126,7 @@ class TestFinalTests(unittest.TestCase):
for line in lines:
if "pickle" in line:
pickle_detectedNum += 1
pickle_accurency = pickle_detectedNum / self.pickle_filesNum
pickle_accurency = pickle_detectedNum / self.pickle_files_num
print(f"pickle files accurency: {pickle_accurency}")
# test pyc files
@ -124,7 +136,7 @@ class TestFinalTests(unittest.TestCase):
for line in lines:
if "pyc" in line:
pyc_detectedNum += 1
pyc_accurency = pyc_detectedNum / self.pyc_filesNum
pyc_accurency = pyc_detectedNum / self.pyc_files_num
print(f"pyc files accurency: {pyc_accurency}")