feat(rules): add AI600-AI900 security rules for agents, RAG, API keys, and output handling

src/pyspector/rules/built-in-rules-ai.toml

@@ -397,3 +397,212 @@ remediation = "Avoid giving LLMs direct SQL execution capabilities. If necessary
 pattern = "create_sql_agent"
 file_pattern = "*.py"
 cwe = "CWE-89"
+
+# -------------------------------------------
+# SECTION: AI600 - Unsafe Agent Behavior & Tool Poisoning
+# -------------------------------------------
+
+[[rule]]
+id = "AI601"
+description = "LLM agent is given unrestricted web browsing capability, risking SSRF and data exfiltration."
+severity = "Critical"
+remediation = "Restrict agent web access to a whitelist of allowed domains. Validate and sanitize all URLs before fetching. Use a proxy or gateway that enforces access policies."
+pattern = "requests\\.get\\s*\\(.*agent|tool.*requests\\.get"
+file_pattern = "*.py"
+cwe = "CWE-918"
+
+[[rule]]
+id = "AI602"
+description = "LLM agent has a tool that executes subprocess calls, risking arbitrary command execution."
+severity = "Critical"
+remediation = "Never give LLM agents direct subprocess execution. If shell access is required, use a sandboxed environment with strict command whitelisting."
+pattern = "subprocess\\.(run|call|Popen|check_output)"
+file_pattern = "*.py"
+cwe = "CWE-78"
+
+[[rule]]
+id = "AI603"
+description = "LLM agent has a tool with unrestricted file write capability, risking arbitrary file overwrite."
+severity = "High"
+remediation = "Constrain file write tools to specific directories. Validate filenames and paths. Implement rate limiting on file operations."
+pattern = "open\\s*\\(.*['\"]w['\"]|['\"]a['\"]"
+file_pattern = "*.py"
+cwe = "CWE-22"
+
+[[rule]]
+id = "AI604"
+description = "Agent tool output is directly concatenated into a prompt without sanitization, enabling indirect prompt injection via tool results."
+severity = "High"
+remediation = "Sanitize and validate all tool outputs before injecting them into prompts. Use delimiters and instruct the LLM to treat tool output as data, not instructions."
+pattern = "f[\"'].*\\{.*tool.*result|\\{.*output.*\\}.*prompt"
+file_pattern = "*.py"
+cwe = "CWE-94"
+
+[[rule]]
+id = "AI605"
+description = "LLM agent is configured with `verbose=True`, which may expose internal reasoning and sensitive data in logs."
+severity = "Medium"
+remediation = "Disable verbose logging in production. If debugging is needed, ensure logs are stored securely and not exposed to end users."
+pattern = "verbose\\s*=\\s*True"
+file_pattern = "*.py"
+cwe = "CWE-200"
+
+[[rule]]
+id = "AI606"
+description = "Agent uses `handle_parsing_errors=False`, which may cause unhandled exceptions to leak internal state."
+severity = "Medium"
+remediation = "Always enable parsing error handling and provide safe fallback responses instead of exposing raw error messages."
+pattern = "handle_parsing_errors\\s*=\\s*False"
+file_pattern = "*.py"
+cwe = "CWE-209"
+
+# -------------------------------------------
+# SECTION: AI700 - RAG Security
+# -------------------------------------------
+
+[[rule]]
+id = "AI701"
+description = "Documents are loaded into a RAG pipeline without content validation, risking embedding poisoning and retrieval manipulation."
+severity = "High"
+remediation = "Validate and sanitize all documents before embedding. Implement content filtering to reject adversarial or malformed inputs."
+pattern = "DirectoryLoader|TextLoader|PDFLoader|UnstructuredFileLoader"
+file_pattern = "*.py"
+cwe = "CWE-345"
+
+[[rule]]
+id = "AI702"
+description = "Vector store similarity threshold is set too low, potentially retrieving irrelevant or adversarial content."
+severity = "Medium"
+remediation = "Set a reasonable similarity threshold (e.g., > 0.7) to filter out low-quality or adversarial retrievals. Monitor retrieval quality metrics."
+pattern = "similarity_search\\s*\\(.*k\\s*=\\s*[0-9]"
+file_pattern = "*.py"
+cwe = "CWE-20"
+
+[[rule]]
+id = "AI703"
+description = "Retrieved context is injected into prompt without size limits, risking context window overflow and DoS."
+severity = "Medium"
+remediation = "Limit the number of retrieved documents and total context length. Implement truncation strategies to stay within model context limits."
+pattern = "combine_docs|stuff_documents_chain"
+file_pattern = "*.py"
+cwe = "CWE-400"
+
+[[rule]]
+id = "AI704"
+description = "Embedding model loaded from an untrusted source can be poisoned to manipulate retrieval results."
+severity = "High"
+remediation = "Use embedding models from trusted, verified sources. Pin model versions and verify checksums when loading from disk."
+pattern = "HuggingFaceEmbeddings|SentenceTransformerEmbeddings"
+file_pattern = "*.py"
+cwe = "CWE-345"
+
+# -------------------------------------------
+# SECTION: AI800 - API Key & Credential Management
+# -------------------------------------------
+
+[[rule]]
+id = "AI801"
+description = "OpenAI API key is hardcoded in the source file."
+severity = "Critical"
+remediation = "Store API keys in environment variables or a secrets manager. Never commit credentials to source control."
+pattern = "openai\\.api_key\\s*=\\s*[\"']sk-"
+file_pattern = "*.py"
+cwe = "CWE-798"
+
+[[rule]]
+id = "AI802"
+description = "Anthropic API key is hardcoded in the source file."
+severity = "Critical"
+remediation = "Store API keys in environment variables or a secrets manager. Never commit credentials to source control."
+pattern = "anthropic\\.api_key\\s*=\\s*[\"']sk-ant-"
+file_pattern = "*.py"
+cwe = "CWE-798"
+
+[[rule]]
+id = "AI803"
+description = "API key is passed as a URL query parameter, risking exposure in logs and referrer headers."
+severity = "High"
+remediation = "Pass API keys in request headers (e.g., Authorization header), never in URL query parameters."
+pattern = "key\\s*=\\s*[\"'].*api.*key|api_key\\s*=.*url"
+file_pattern = "*.py"
+cwe = "CWE-598"
+
+[[rule]]
+id = "AI804"
+description = "Cohere API key is hardcoded in the source file."
+severity = "Critical"
+remediation = "Store API keys in environment variables or a secrets manager. Never commit credentials to source control."
+pattern = "cohere\\.Client\\s*\\(.*api_key\\s*=\\s*[\"']"
+file_pattern = "*.py"
+cwe = "CWE-798"
+
+# -------------------------------------------
+# SECTION: AI900 - Output Handling & DoS
+# -------------------------------------------
+
+[[rule]]
+id = "AI901"
+description = "Unsafe YAML parsing of LLM output can lead to arbitrary object instantiation and RCE."
+severity = "Critical"
+remediation = "Use yaml.safe_load() or yaml.safe_loads() instead of yaml.load() with untrusted LLM output."
+pattern = "yaml\\.load\\s*\\("
+exclude_pattern = "yaml\\.safe_load"
+file_pattern = "*.py"
+cwe = "CWE-502"
+
+[[rule]]
+id = "AI902"
+description = "JSON parsing of LLM output without size limits can lead to memory exhaustion DoS."
+severity = "Medium"
+remediation = "Limit the maximum size of LLM output before parsing. Use streaming JSON parsers for large responses."
+pattern = "json\\.loads\\s*\\("
+file_pattern = "*.py"
+cwe = "CWE-400"
+
+[[rule]]
+id = "AI903"
+description = "LLM output is passed directly to exec() or eval(), risking arbitrary code execution."
+severity = "Critical"
+remediation = "Never execute LLM-generated code without sandboxing. Use ast.literal_eval() for data structures or a restricted execution environment."
+pattern = "(exec|eval)\\s*\\(.*llm|.*response|.*output|.*completion"
+file_pattern = "*.py"
+cwe = "CWE-94"
+
+[[rule]]
+id = "AI904"
+description = "LLM response is directly rendered as HTML without sanitization, risking XSS."
+severity = "High"
+remediation = "Sanitize LLM output before rendering as HTML. Use a markup sanitizer like bleach to strip dangerous tags and attributes."
+pattern = "innerHTML\\s*=|dangerouslySetInnerHTML|render_template_string.*llm|.*response"
+file_pattern = "*.py"
+cwe = "CWE-79"
+
+# -------------------------------------------
+# NEW TAINT SOURCES & SINKS for AI600-AI900
+# -------------------------------------------
+
+[[taint_source]]
+id = "AITS11"
+description = "Data retrieved from a vector store in a RAG pipeline is considered tainted."
+function_call = "langchain_community.vectorstores Chroma.similarity_search"
+taint_target = "return"
+
+[[taint_source]]
+id = "AITS12"
+description = "Output from a web scraping tool used by an agent is considered tainted."
+function_call = "requests.get"
+taint_target = "return"
+
+[[taint_sink]]
+id = "AISK11"
+vulnerability_id = "AI601"
+description = "Tainted data is used as a URL in an agent's web browsing tool."
+function_call = "requests.get"
+vulnerable_parameter_index = 0
+
+[[taint_sink]]
+id = "AISK12"
+vulnerability_id = "AI901"
+description = "Tainted data is passed to yaml.load for deserialization."
+function_call = "yaml.load"
+vulnerable_parameter_index = 0

tests/unit/test_ai_rules.py

@@ -208,3 +208,167 @@ def test_exclude_pattern_suppresses_safe_or_comment_cases(self, code):
         rule = _ai_rule("AI202")
         assert re.search(rule["pattern"], code)
         assert re.search(rule["exclude_pattern"], code)
+
+
+# -------------------------------------------
+# Tests for AI600 - Unsafe Agent Behavior
+# -------------------------------------------
+
+class TestAI600AgentBehavior:
+    def test_ai601_metadata(self):
+        rule = _ai_rule("AI601")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-918"
+
+    def test_ai601_pattern_matches(self):
+        rule = _ai_rule("AI601")
+        assert re.search(rule["pattern"], "response = requests.get(url, headers=headers)")
+
+    def test_ai602_metadata(self):
+        rule = _ai_rule("AI602")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-78"
+
+    def test_ai602_pattern_matches(self):
+        rule = _ai_rule("AI602")
+        assert re.search(rule["pattern"], "subprocess.run(command, shell=True)")
+
+    def test_ai603_metadata(self):
+        rule = _ai_rule("AI603")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-22"
+
+    def test_ai604_metadata(self):
+        rule = _ai_rule("AI604")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-94"
+
+    def test_ai605_metadata(self):
+        rule = _ai_rule("AI605")
+        assert rule["severity"] == "Medium"
+        assert rule["cwe"] == "CWE-200"
+
+    def test_ai605_pattern_matches(self):
+        rule = _ai_rule("AI605")
+        assert re.search(rule["pattern"], "agent = initialize_agent(verbose=True)")
+
+    def test_ai606_metadata(self):
+        rule = _ai_rule("AI606")
+        assert rule["severity"] == "Medium"
+        assert rule["cwe"] == "CWE-209"
+
+    def test_ai606_pattern_matches(self):
+        rule = _ai_rule("AI606")
+        assert re.search(rule["pattern"], "agent = initialize_agent(handle_parsing_errors=False)")
+
+
+# -------------------------------------------
+# Tests for AI700 - RAG Security
+# -------------------------------------------
+
+class TestAI700RAGSecurity:
+    def test_ai701_metadata(self):
+        rule = _ai_rule("AI701")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-345"
+
+    def test_ai701_pattern_matches(self):
+        rule = _ai_rule("AI701")
+        assert re.search(rule["pattern"], "loader = DirectoryLoader('./docs')")
+
+    def test_ai702_metadata(self):
+        rule = _ai_rule("AI702")
+        assert rule["severity"] == "Medium"
+        assert rule["cwe"] == "CWE-20"
+
+    def test_ai703_metadata(self):
+        rule = _ai_rule("AI703")
+        assert rule["severity"] == "Medium"
+        assert rule["cwe"] == "CWE-400"
+
+    def test_ai704_metadata(self):
+        rule = _ai_rule("AI704")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-345"
+
+    def test_ai704_pattern_matches(self):
+        rule = _ai_rule("AI704")
+        assert re.search(rule["pattern"], "embeddings = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')")
+
+
+# -------------------------------------------
+# Tests for AI800 - API Key Management
+# -------------------------------------------
+
+class TestAI800APIKeyManagement:
+    def test_ai801_metadata(self):
+        rule = _ai_rule("AI801")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-798"
+
+    def test_ai801_pattern_matches(self):
+        rule = _ai_rule("AI801")
+        assert re.search(rule["pattern"], 'openai.api_key = "sk-abc123def456"')
+
+    def test_ai801_pattern_no_match_env_var(self):
+        rule = _ai_rule("AI801")
+        assert not re.search(rule["pattern"], 'openai.api_key = os.getenv("OPENAI_API_KEY")')
+
+    def test_ai802_metadata(self):
+        rule = _ai_rule("AI802")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-798"
+
+    def test_ai802_pattern_matches(self):
+        rule = _ai_rule("AI802")
+        assert re.search(rule["pattern"], 'anthropic.api_key = "sk-ant-abc123"')
+
+    def test_ai803_metadata(self):
+        rule = _ai_rule("AI803")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-598"
+
+    def test_ai804_metadata(self):
+        rule = _ai_rule("AI804")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-798"
+
+    def test_ai804_pattern_matches(self):
+        rule = _ai_rule("AI804")
+        assert re.search(rule["pattern"], "cohere.Client(api_key='abc123')")
+
+
+# -------------------------------------------
+# Tests for AI900 - Output Handling & DoS
+# -------------------------------------------
+
+class TestAI900OutputHandling:
+    def test_ai901_metadata(self):
+        rule = _ai_rule("AI901")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-502"
+
+    def test_ai901_pattern_matches(self):
+        rule = _ai_rule("AI901")
+        assert re.search(rule["pattern"], "data = yaml.load(llm_output)")
+
+    def test_ai901_excludes_safe_load(self):
+        rule = _ai_rule("AI901")
+        code = "data = yaml.safe_load(llm_output)"
+        assert re.search(rule["pattern"], code)
+        assert re.search(rule["exclude_pattern"], code)
+
+    def test_ai902_metadata(self):
+        rule = _ai_rule("AI902")
+        assert rule["severity"] == "Medium"
+        assert rule["cwe"] == "CWE-400"
+
+    def test_ai903_metadata(self):
+        rule = _ai_rule("AI903")
+        assert rule["severity"] == "Critical"
+        assert rule["cwe"] == "CWE-94"
+
+    def test_ai904_metadata(self):
+        rule = _ai_rule("AI904")
+        assert rule["severity"] == "High"
+        assert rule["cwe"] == "CWE-79"