74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
import unittest
|
|
import shutil
|
|
|
|
from detection.utils import read_file_content
|
|
from .final_tests_util import clone_repo, Path, inject_random_backdoor
|
|
from detection.Regexdetection import find_dangerous_functions
|
|
from detection.GPTdetection import detectGPT
|
|
|
|
|
|
class TestFinalTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.path = "./tmp/repo"
|
|
shutil.rmtree(self.path, ignore_errors=True)
|
|
clone_repo("https://github.com/injetlee/Python.git", self.path)
|
|
sampleRate = 0.1
|
|
self.inject_reslt = inject_random_backdoor(self.path, sample_rate=sampleRate)
|
|
self.injectedNum = len(self.inject_reslt)
|
|
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.trueRate = self.injectedNum / self.py_filesNum
|
|
print(self.trueRate)
|
|
|
|
# test backdoor code in python files
|
|
def test_final_tests_pycode(self):
|
|
detectedNum = 0
|
|
possibly_dangerous_file = []
|
|
for file in self.all_python_files:
|
|
content = read_file_content(str(file))
|
|
results = find_dangerous_functions(content, ".py")
|
|
if (
|
|
len(results["high"]) > 0
|
|
or len(results["medium"]) > 0
|
|
or len(results["low"]) > 0
|
|
):
|
|
detectedNum += 1
|
|
possibly_dangerous_file.append(file)
|
|
print(detectedNum / self.py_filesNum)
|
|
self.assertAlmostEqual(detectedNum, self.py_filesNum, places=1)
|
|
GPTdetectedNum = 0
|
|
|
|
for i in possibly_dangerous_file:
|
|
content = read_file_content(str(i))
|
|
results = {}
|
|
try:
|
|
results = detectGPT(content)
|
|
if (
|
|
len(results["high"]) > 0
|
|
or len(results["medium"]) > 0
|
|
or len(results["low"]) > 0
|
|
):
|
|
GPTdetectedNum += 1
|
|
print(GPTdetectedNum)
|
|
self.assertGreaterEqual(GPTdetectedNum, detectedNum)
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
# test pickle files
|
|
pickle_detectedNum = 0
|
|
pickle_tureNum = len(list(Path(self.path).glob("*.pickle")))
|
|
|
|
self.assertAlmostEqual(pickle_detectedNum, pickle_tureNum, places=1)
|
|
|
|
# test pyc files
|
|
pyc_detectedNum = 0
|
|
pyc_tureNum = len(list(Path(self.path).glob("*.pyc")))
|
|
self.assertAlmostEqual(pyc_detectedNum, pyc_tureNum, places=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|