feat: update test case

This commit is contained in:
sangge-redmi 2024-06-04 18:11:14 +08:00
parent fbeba5b4fc
commit 6e1c0e5ae6
2 changed files with 56 additions and 32 deletions

View File

@ -3,9 +3,8 @@ from git import Repo # type: ignore
import random
from pathlib import Path
import pickle
import marshal
import importlib.util
import os
import py_compile
def clone_repo(repo_url: str, clone_dir: str) -> None:
@ -175,20 +174,13 @@ def inject_pyc_backdoor(root_path: str) -> None:
for path in paths:
backdoor_id = random.randrange(0, len(backdoors))
backdoor = backdoors[backdoor_id]
filename = os.path.join(path, f"backdoor{backdoor_id}.pyc")
py_filename = os.path.join(path, f"backdoor{backdoor_id}.py")
pyc_filename = os.path.join(path, f"backdoor{backdoor_id}.pyc")
with open(py_filename, "w") as f:
f.write(backdoor)
# Compile the string to a code object
code = compile(backdoor, filename, "exec")
# Create a code object header
header = importlib.util.MAGIC_NUMBER
if hasattr(importlib.util, "SOURCE_SUFFIXES"):
header += b"\x00" * 4
# Write the .pyc file
with open(filename, "wb") as file:
file.write(header)
marshal.dump(code, file)
py_compile.compile(py_filename, cfile=pyc_filename)
os.remove(py_filename)
if __name__ == "__main__":

View File

@ -1,5 +1,6 @@
import unittest
import shutil
import os
from detection.utils import read_file_content
from .final_tests_util import clone_repo, Path, inject_random_backdoor
@ -9,21 +10,34 @@ from detection.GPTdetection import detectGPT
class TestFinalTests(unittest.TestCase):
def setUp(self) -> None:
self.path = "./tmp/repo"
self.path = "./tmp/repo/"
shutil.rmtree(self.path, ignore_errors=True)
clone_repo("https://github.com/injetlee/Python.git", self.path)
if not os.path.exists("/tmp/Python/"):
clone_repo("https://github.com/TheAlgorithms/Python.git", "/tmp/Python")
shutil.copytree("/tmp/Python", self.path)
sampleRate = 0.1
self.inject_reslt = inject_random_backdoor(self.path, sample_rate=sampleRate)
self.injectedNum = len(self.inject_reslt)
self.inject_result = inject_random_backdoor(
self.path, sample_rate=sampleRate, pyc=True, pickle=True
)
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.trueRate = self.injectedNum / self.py_filesNum
print(self.trueRate)
# test backdoor code in python files
all_pickle_files = list(project_path.rglob("*.pickle"))
self.pickle_filesNum = len(all_pickle_files)
all_pyc_files = list(project_path.rglob("*.pyc"))
self.pyc_filesNum = len(all_pyc_files)
os.system(
"python -m detection " + self.path + " -o " + self.path + "output.txt"
)
def test_final_tests_pycode(self):
# test backdoor code in python files
detectedNum = 0
possibly_dangerous_file = []
for file in self.all_python_files:
@ -37,7 +51,6 @@ class TestFinalTests(unittest.TestCase):
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:
@ -52,21 +65,40 @@ class TestFinalTests(unittest.TestCase):
):
GPTdetectedNum += 1
print(GPTdetectedNum)
self.assertGreaterEqual(GPTdetectedNum, detectedNum)
except Exception as e:
print(e)
# print(e)
pass
# test injected code
with open(self.path + "output.txt", "r") as f:
lines = f.readlines()
injected_detectedNum = 0
for line in lines:
if "py:" in line:
injected_detectedNum += 1
injected_accurency = injected_detectedNum / self.injectedNum
print(f"injected files accurency: {injected_accurency}")
# test pickle files
with open(self.path + "output.txt", "r") as f:
lines = f.readlines()
pickle_detectedNum = 0
pickle_tureNum = len(list(Path(self.path).glob("*.pickle")))
self.assertAlmostEqual(pickle_detectedNum, pickle_tureNum, places=1)
for line in lines:
if "pickle" in line:
pickle_detectedNum += 1
pickle_accurency = pickle_detectedNum / self.pickle_filesNum
print(f"pickle files accurency: {pickle_accurency}")
# test pyc files
with open(self.path + "output.txt", "r") as f:
lines = f.readlines()
pyc_detectedNum = 0
pyc_tureNum = len(list(Path(self.path).glob("*.pyc")))
self.assertAlmostEqual(pyc_detectedNum, pyc_tureNum, places=1)
for line in lines:
if "pyc" in line:
pyc_detectedNum += 1
pyc_accurency = pyc_detectedNum / self.pyc_filesNum
print(f"pyc files accurency: {pyc_accurency}")
if __name__ == "__main__":