From d960bfa74328f210242efe0b6b15fbaed37abdd2 Mon Sep 17 00:00:00 2001 From: Mustafa Senoglu Date: Mon, 3 Aug 2026 14:46:30 +0300 Subject: [PATCH] feat(rules): add AI600-AI900 security rules for agent behavior, RAG, API keys, and output handling Adds 18 new rules covering: - AI600: Unsafe agent behavior & tool poisoning (web browsing, subprocess, file write, indirect injection) - AI700: RAG security (embedding poisoning, context overflow, untrusted sources) - AI800: API key management (OpenAI, Anthropic, Cohere hardcoded keys) - AI900: Output handling & DoS (YAML unsafe load, JSON DoS, exec/eval of LLM output, XSS) Also adds 4 new taint sources/sinks for RAG and agent web tool flows. Closes #91 --- src/pyspector/rules/built-in-rules-ai.toml | 209 +++++++++++++++++++++ tests/unit/test_ai_rules.py | 164 ++++++++++++++++ 2 files changed, 373 insertions(+) diff --git a/src/pyspector/rules/built-in-rules-ai.toml b/src/pyspector/rules/built-in-rules-ai.toml index 7a37ca0..d6a15b1 100644 --- a/src/pyspector/rules/built-in-rules-ai.toml +++ b/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 diff --git a/tests/unit/test_ai_rules.py b/tests/unit/test_ai_rules.py index b56cc10..3cdeefd 100644 --- a/tests/unit/test_ai_rules.py +++ b/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"