Merge pull request 'tests/final-tests 完成最终代码' (#34) from tests/final-tests into main

Reviewed-on: #34
Reviewed-by: dqy <dqy@noreply.localhost>

detection/GPTdetection.py

@@ -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

detection/__main__.py

@@ -21,7 +21,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__",
@@ -111,6 +111,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)
@@ -388,13 +389,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
-                if file_extension in [".pkl", ".pickle"]:
+                # 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), "result": res})
                     continue

tests/final_tests_util.py

@@ -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_id = random.randrange(0, len(backdoors_pickle))
-        backdoor = backdoors[backdoor_id]
+        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)