detection/GPTdetection.py

@@ -7,8 +7,12 @@ import openai
 
 
 def detect_gpt(filename: str):
+    content = read_file_content(filename)
+    return detectGPT(content)
+
+def detectGPT(content: str):
     client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-    text = read_file_content(filename)
+    text = content
     # client = openai.OpenAI(api_key="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") #测试用key
     response = client.chat.completions.create(
         messages=[

tests/python-test.yml

@@ -16,3 +16,6 @@ jobs:
         run: pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
       - name: Run tests
         run: python -m unittest discover -s tests
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          http_proxy: http://192.168.1.3:10809

tests/test_backdoor_detection.py

@@ -1,6 +1,7 @@
 import unittest
 
 from detection.backdoor_detection import find_dangerous_functions
+from detection.GPTdetection import *
 
 
 class TestBackdoorDetection(unittest.TestCase):
@@ -55,6 +56,25 @@ class TestBackdoorDetection(unittest.TestCase):
             results["medium"],
         )
 
+    def test_gpt_risk_detection(self):
+        content = """import os
+        os.system('rm -rf /')   # high risk
+        exec('print("Hello")')  # high risk
+        eval('2 + 2')   # high risk
+        """
+        results = detectGPT(content)
+        self.assertEqual(len(results["high"]), 3)
+
+    def test_gpt_no_risk_detection(self):
+        content = """a = 10
+        b = a + 5
+        print('This should not be detected as risky.')
+        """
+        results = detectGPT(content)
+        self.assertEqual(len(results["high"]), 0)
+        self.assertEqual(len(results["medium"]), 0)
+        self.assertEqual(len(results["low"]), 0)
+
 
 if __name__ == "__main__":
     unittest.main()