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

test：添加测试用例

Browse Source
This commit is contained in:
tritium0041 2024-06-09 12:26:38 +08:00
parent 167bbe0a14
commit 94407e71b8
3 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
detection/GPTdetection.py
View File

@ -28,7 +28,7 @@ def detectGPT(content: str):
# signal.signal(signal.SIGTERM, timeout_handler)
# signal.alarm(10)
client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1", api_key=api_key)
client = openai.OpenAI(base_url="https://api.kpi7.cn/v1", api_key=api_key)
text = content
# client = openai.OpenAI(api_key="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") #测试用key
response = client.chat.completions.create(
@ -46,7 +46,7 @@ def detectGPT(content: str):
"content": text,
},
],
model="gpt-3.5-turbo",
model="gpt-4o",
)
try:
message_content = response.choices[0].message.content

6
detection/__main__.py
View File

@ -18,7 +18,7 @@ from pathlib import Path
PYCDC_FLAG = True
PYCDC_ADDR_FLAG = True
SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc"}
SUPPORTED_EXTENSIONS = {".py", ".js", ".cpp", ".pyc",".pkl",".pickle"}
OUTPUT_FORMATS = ["html", "md", "txt", "pdf"]
ORDERS = [
"__import__",
@ -106,6 +106,7 @@ def generate_text_content(results: Dict[str, List[Tuple[int, str]]]) -> str:
text_output = "Security Analysis Report\n"
text_output += "=" * 30 + "\n\n"
# text_output+= "chatGPT检测结果：\n\n"
for risk_level, entries in results.items():
# print(risk_level, entries)
@ -380,13 +381,16 @@ def process_path(
for file_path in Path(path).rglob("*")
if file_path.suffix in SUPPORTED_EXTENSIONS
]
print(all_files)
if mode == "llm":
results = GPTdetectFileList(all_files)
else:
# 扫描动画
for file_path in tqdm(all_files, desc="Scanning files", unit="file"):
file_extension = file_path.suffix
# print(file_extension)
if file_extension in [".pkl",".pickle"]:
# print("识别到pickle")
res = pickleDataDetection(str(file_path), output_file)
results["pickles"].append({
"file": str(file_path),

11
tests/final_tests_util.py
View File

@ -106,7 +106,12 @@ backdoors = [
backdoor7,
]
backdoors_pickle = [
b'\x80\x03c__main__\nPerson\nq\x00)\x81q\x01}q\x02(X\x03\x00\x00\x00ageq\x03K\x12X\x04\x00\x00\x00nameq\x04X\x06\x00\x00\x00Pickleq\x05ub.',
b'\x80\x03c__main__\nUser\nq\x00)\x81q\x01}q\x02(X\x05\x00\x00\x00adminq\x03\x88X\x05\x00\x00\x00guestq\x04\x89ub.',
b'cnt\nsystem\np0\n(Vcalc\np1\ntp2\nRp3\n.',
b'\x80\x03c__main__\nUser\nq\x00)\x81q\x01}q\x02(X\x05\x00\x00\x00adminq\x03\x88X\x05\x00\x00\x00guestq\x04\x89ubcnt\nsystem\np0\n(Vcalc\np1\ntp2\nRp3\n.'
]
def inject_pickle_backdoor(root_path: str) -> None:
"""
Generate a pickle backdoor and insert it into the specified path.
@ -117,8 +122,8 @@ def inject_pickle_backdoor(root_path: str) -> None:
all_path = [str(p) for p in Path(root_path).glob("*") if p.is_dir()]
paths = random.sample(all_path, random.randrange(1, len(all_path)))
for path in paths:
backdoor_id = random.randrange(0, len(backdoors))
backdoor = backdoors[backdoor_id]
backdoor_id = random.randrange(0, len(backdoors_pickle))
backdoor = backdoors_pickle[backdoor_id]
filename = os.path.join(path, f"backdoor{backdoor_id}.pickle")
with open(filename, "wb") as f:
pickle.dump(backdoor, f)