detection/cngptdetection.py

@@ -0,0 +1,97 @@
+import json
+import requests
+import signal
+from typing import Dict, List, Tuple     # 用于类型提示的模块，使用了 Dict, List, Tuple 进行类型注解。
+
+# 参考文档：https://blog.csdn.net/weixin_73654895/article/details/133799269
+
+class TimeoutException(Exception):
+    """Custom exception to handle timeouts."""
+    pass
+
+
+def timeout_handler(signum, frame):
+    """Handle the SIGALRM signal by raising a TimeoutException."""
+    raise TimeoutException
+
+
+def get_baidu_access_token(api_key: str, secret_key: str) -> str:
+    """
+    Retrieve the access token from Baidu API using API key and Secret key.
+
+    Args:
+        api_key (str): The API key for Baidu API.
+        secret_key (str): The Secret key for Baidu API.
+
+    Returns:
+        str: The access token.
+    """
+    url = "https://aip.baidubce.com/oauth/2.0/token"
+    params = {"grant_type": "client_credentials", "client_id": api_key, "client_secret": secret_key}
+    response = requests.post(url, params=params)
+    response_data = response.json()
+    if 'access_token' not in response_data:
+        raise ValueError("Error: Could not retrieve access token.")
+    return str(response_data["access_token"])
+
+
+def cndetectGPT(content: str) -> Dict[str, List[Tuple[int, str]]]:
+    """
+    Detect potential security vulnerabilities in the provided code content using Baidu's AI model.
+
+    Args:
+        content (str): The code content to be analyzed.
+
+    Returns:
+        Dict[str, List[Tuple[int, str]]]: Classified results of detected vulnerabilities.
+    """
+    API_KEY = "DUBWNIrB6QJLOsLkpnEz2ZZa"
+    SECRET_KEY = "9WK4HIV2n9r1ePPirqD4EQ6Ea33rH1m7"
+
+    # Set alarm timer
+    signal.signal(signal.SIGTERM, timeout_handler)
+    signal.alarm(10)
+
+    try:
+        access_token = get_baidu_access_token(API_KEY, SECRET_KEY)
+        url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token={access_token}"
+        
+        payload = json.dumps({
+            "messages": [
+                {
+                    "role": "system",
+                    "content": "You are a Python code reviewer. Read the code below and identify any potential security vulnerabilities. Classify them by risk level (high, medium, low, none). Only report the line number and the risk level.\nYou should output the result as json format in one line. For example: "
+                               '[{"Line": {the line number}, "Risk": "{choose from (high,medium,low)}","Reason":"{how it is vulnerable}"}] Each of these three fields is required.\n'
+                               "You are required to only output the json format. Do not output any other information.\n"
+                },
+                {
+                    "role": "user",
+                    "content": content
+                }
+            ]
+        })
+        
+        headers = {
+            'Content-Type': 'application/json'
+        }
+        
+        response = requests.post(url, headers=headers, data=payload)
+        response_data = response.json()
+        message_content = response_data.get('result', None)
+        if message_content is None:
+            raise ValueError("API response content is None")
+        res_json = json.loads(message_content)
+        
+    except json.JSONDecodeError:
+        raise ValueError("Error: Could not parse the response. Please try again.")
+    except TimeoutException:
+        raise TimeoutException("The API call timed out")
+    finally:
+        signal.alarm(0)
+
+    classified_results = {"high": [], "medium": [], "low": [], "none": []}
+    for res in res_json:
+        classified_results[res["Risk"]].append(
+            (res["Line"], content.split("\n")[res["Line"] - 1].strip())
+        )
+    return classified_results

tests/test_CN_GPT_detection.py

@@ -0,0 +1,35 @@
+import unittest
+import warnings
+import os
+from detection.cngptdetection import cndetectGPT
+
+class TestBackdoorDetection(unittest.TestCase):
+    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 = cndetectGPT(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 = cndetectGPT(content)
+        self.assertEqual(len(results["high"]), 0)
+        self.assertEqual(len(results["medium"]), 0)
+        self.assertEqual(len(results["low"]), 0)
+
+    def test_gpt_env_no_set(self):
+        if os.getenv("BAIDU_API_KEY") is None or os.getenv("BAIDU_SECRET_KEY") is None:
+            self.skipTest("BAIDU_API_KEY or BAIDU_SECRET_KEY is not set")
+        content = "print('test test')"
+        with self.assertRaises(ValueError):
+            cndetectGPT(content)
+
+
+if __name__ == "__main__":
+    unittest.main()