Merge branch 'main' into feature/pickle-data

tests/test_CN_GPT_detection.py

@@ -0,0 +1,40 @@
+import unittest
+import warnings
+import os
+import json
+
+from detection.cngptdetection import detectGPT
+
+class TestBackdoorDetection(unittest.TestCase):
+    def test_gpt_risk_detection(self):
+        if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
+            warnings.warn("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set, test skipped.", UserWarning)
+            self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
+
+        content = """import os
+        os.system('rm -rf /')   # high risk
+        exec('print("Hello")')  # high risk
+        eval('2 + 2')   # high risk
+        """
+        results1 = detectGPT(content)
+        classified_results = json.loads(results1)
+        self.assertEqual(len(classified_results["high"]), 3)
+
+    def test_gpt_no_risk_detection(self):
+        if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
+            warnings.warn("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set, test skipped.", UserWarning)
+            self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
+
+        content = """a = 10
+        b = a + 5
+        print('This should not be detected as risky.')
+        """
+        results2 = detectGPT(content)
+        classified_results = json.loads(results2)
+        self.assertEqual(len(classified_results["high"]), 0)
+        self.assertEqual(len(classified_results["medium"]), 0)
+        self.assertEqual(len(classified_results["low"]), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_backdoor_detection.py

@@ -1,7 +1,7 @@
 import unittest
 import warnings
 
-from detection.backdoor_detection import find_dangerous_functions
+from detection.__main__ import find_dangerous_functions
 from detection.GPTdetection import detectGPT
 import os
 
@@ -90,6 +90,23 @@ class TestBackdoorDetection(unittest.TestCase):
         with self.assertRaises(ValueError):
             detectGPT(content)
 
+    def test_find_dangerous_functions_pyc(self):
+        file_content = """import os
+        os.system('rm -rf /')
+        """
+        file_extension = ".pyc"
+
+        expected_result = {
+            "high": [(2, "os.system('rm -rf /')")],
+            "medium": [],
+            "low": [],
+            "none": [],
+        }
+
+        result = find_dangerous_functions(file_content, file_extension)
+
+        self.assertEqual(result, expected_result)
+
 
 if __name__ == "__main__":
     unittest.main()