From d960bfa74328f210242efe0b6b15fbaed37abdd2 Mon Sep 17 00:00:00 2001 From: Mustafa Senoglu Date: Mon, 3 Aug 2026 14:46:30 +0300 Subject: [PATCH 1/2] 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" From 2a73c9910b039d5b6b0c95449e817cb659751a42 Mon Sep 17 00:00:00 2001 From: Mustafa Senoglu Date: Fri, 14 Aug 2026 17:25:04 +0300 Subject: [PATCH 2/2] fix(rules): address review feedback on AI600-AI900 rules - Remove AI601/AI603/AI604: line-based regex cannot establish agent context; SSRF and file access already covered by AI501/AISK08 and AI502/AISK09 taint analysis - Remove AITS12/AISK11: AITS12 tainted every requests.get response, AISK11 duplicated the existing AISK08 sink - AI602: require shell=True instead of flagging all subprocess usage - AI702: pattern now checks score_threshold (was checking top-k count) - AI902/AI903/AI904: narrow patterns to LLM-output variables; drop broad .*response|.*output|.*completion alternatives that matched any line containing those words - AI904: exclude sanitized output (sanitize|bleach) - AITS11: fix malformed function path, match instance calls via .similarity_search - Add negative tests: plain HTTP clients, subprocess without shell, plain json.loads, other embedding models, sanitized HTML, top-k without threshold, plain eval, multiline taint flow - Taint tests: trusted constants stay untainted, no duplicate AI501/AI601 findings --- src/pyspector/rules/built-in-rules-ai.toml | 77 ++++----- tests/unit/test_ai_rules.py | 180 ++++++++++++++++++--- 2 files changed, 186 insertions(+), 71 deletions(-) diff --git a/src/pyspector/rules/built-in-rules-ai.toml b/src/pyspector/rules/built-in-rules-ai.toml index d6a15b1..3cdd76c 100644 --- a/src/pyspector/rules/built-in-rules-ai.toml +++ b/src/pyspector/rules/built-in-rules-ai.toml @@ -252,6 +252,7 @@ description = "Use of joblib.load can be insecure for untrusted model files." severity = "High" remediation = "Joblib can use pickle under the hood. Treat .joblib files as potentially malicious and only load from trusted sources." pattern = "joblib\\.load" +exclude_pattern = "^\\s*#" file_pattern = "*.py" cwe = "CWE-502" @@ -402,46 +403,37 @@ 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" +# AI601 (agent web browsing / SSRF) was removed: its regex could not reliably +# match multiline or separately-defined agent tools, and its taint sink +# (requests.get) duplicated the existing AISK08 sink. Agent-side SSRF is +# already covered by AI501/AISK08 taint analysis. [[rule]] id = "AI602" -description = "LLM agent has a tool that executes subprocess calls, risking arbitrary command execution." +description = "LLM agent has a tool that executes subprocess calls with shell=True, 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)" +confidence = "Low" +remediation = "Never give LLM agents direct subprocess execution. If shell access is required, use a sandboxed environment with strict command whitelisting. Prefer shell=False with explicit argument lists." +pattern = "subprocess\\.(?:run|call|Popen|check_output)\\s*\\([^)]*shell\\s*=\\s*True" 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" +# AI603 (unrestricted file write in agent tools) was removed: its pattern had +# an ungrouped alternation that matched unrelated code (e.g. mode = "a"), and +# ordinary file writes cannot be attributed to an agent tool via line-based +# regex. Agent-side filesystem access is already covered by AI502/AISK09 +# taint analysis. -[[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" +# AI604 (tool output concatenated into prompts) was removed: detecting +# unsanitized tool output requires data-flow analysis. The equivalent flow +# (LLM/tool output reaching a prompt template) is covered by AITS04 -> AISK01 +# (AI101) taint analysis. [[rule]] id = "AI605" description = "LLM agent is configured with `verbose=True`, which may expose internal reasoning and sensitive data in logs." severity = "Medium" +confidence = "Low" 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" @@ -464,6 +456,7 @@ cwe = "CWE-209" id = "AI701" description = "Documents are loaded into a RAG pipeline without content validation, risking embedding poisoning and retrieval manipulation." severity = "High" +confidence = "Low" 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" @@ -471,10 +464,11 @@ cwe = "CWE-345" [[rule]] id = "AI702" -description = "Vector store similarity threshold is set too low, potentially retrieving irrelevant or adversarial content." +description = "Vector store similarity search uses an explicit score_threshold; ensure the threshold is high enough to filter out low-quality or adversarial retrievals." severity = "Medium" +confidence = "Low" 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]" +pattern = "similarity_search\\s*\\([^)]*score_threshold\\s*=\\s*0\\.[0-9]+" file_pattern = "*.py" cwe = "CWE-20" @@ -482,6 +476,7 @@ cwe = "CWE-20" id = "AI703" description = "Retrieved context is injected into prompt without size limits, risking context window overflow and DoS." severity = "Medium" +confidence = "Low" 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" @@ -491,6 +486,7 @@ cwe = "CWE-400" id = "AI704" description = "Embedding model loaded from an untrusted source can be poisoned to manipulate retrieval results." severity = "High" +confidence = "Low" remediation = "Use embedding models from trusted, verified sources. Pin model versions and verify checksums when loading from disk." pattern = "HuggingFaceEmbeddings|SentenceTransformerEmbeddings" file_pattern = "*.py" @@ -554,8 +550,9 @@ cwe = "CWE-502" id = "AI902" description = "JSON parsing of LLM output without size limits can lead to memory exhaustion DoS." severity = "Medium" +confidence = "Low" remediation = "Limit the maximum size of LLM output before parsing. Use streaming JSON parsers for large responses." -pattern = "json\\.loads\\s*\\(" +pattern = "json\\.loads\\s*\\(\\s*\\w*(?:output|response|completion)\\w*\\s*\\)" file_pattern = "*.py" cwe = "CWE-400" @@ -564,7 +561,7 @@ 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" +pattern = "(?:exec|eval)\\s*\\(\\s*\\w*(?:llm|response|output|completion)\\w*\\s*\\)" file_pattern = "*.py" cwe = "CWE-94" @@ -573,7 +570,8 @@ 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" +pattern = "innerHTML\\s*=|dangerouslySetInnerHTML|render_template_string\\s*\\([^)]*llm" +exclude_pattern = "sanitize|bleach" file_pattern = "*.py" cwe = "CWE-79" @@ -584,22 +582,9 @@ cwe = "CWE-79" [[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" +function_call = ".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" diff --git a/tests/unit/test_ai_rules.py b/tests/unit/test_ai_rules.py index 3cdeefd..fbe2792 100644 --- a/tests/unit/test_ai_rules.py +++ b/tests/unit/test_ai_rules.py @@ -166,16 +166,15 @@ def test_joblib_model_load_metadata(self): assert rule["pattern"] == r"joblib\.load" def test_joblib_model_load_fires(self): - code = """ - model = joblib.load(model_path) - """ - assert fires(code, "AI204") + rule = _ai_rule("AI204") + assert re.search(rule["pattern"], "model = joblib.load(model_path)") + assert not re.search(rule["exclude_pattern"], "model = joblib.load(model_path)") def test_commented_joblib_model_load_safe(self): - code = """ - # model = joblib.load(model_path) - """ - assert not fires(code, "AI204") + rule = _ai_rule("AI204") + code = "# model = joblib.load(model_path)" + assert re.search(rule["pattern"], code) + assert re.search(rule["exclude_pattern"], code) class TestAI202: def test_rule_metadata(self): rule = _ai_rule("AI202") @@ -215,33 +214,78 @@ def test_exclude_pattern_suppresses_safe_or_comment_cases(self, code): # ------------------------------------------- class TestAI600AgentBehavior: - def test_ai601_metadata(self): - rule = _ai_rule("AI601") - assert rule["severity"] == "Critical" - assert rule["cwe"] == "CWE-918" + def test_ai601_removed_covered_by_ai501_taint(self): + rule_ids = [rule["id"] for rule in toml.loads(RULES_PATH.read_text(encoding="utf-8"))["rule"]] + assert "AI601" not in rule_ids + + def test_ai601_removed_taint_flow_fires_ai501_not_ai601(self): + code = """ + url = input("target: ") + requests.get(url) + """ + results = run_pyspector_ai(code) + assert any(result["rule_id"] == "AI501" for result in results) + assert not any(result["rule_id"] == "AI601" for result in results) + + def test_ai601_removed_trusted_url_not_tainted(self): + code = """ + response = requests.get("https://api.example.com/data") + """ + assert not fires(code, "AI501") + + def test_ai501_multiline_requests_get_taint(self): + code = """ + url = input("target: ") + requests.get( + url, + headers={"User-Agent": "agent"}, + ) + """ + assert fires(code, "AI501") - def test_ai601_pattern_matches(self): - rule = _ai_rule("AI601") - assert re.search(rule["pattern"], "response = requests.get(url, headers=headers)") + def test_ai501_tainted_url_from_request_body(self): + code = """ + url = request.form.get("url") + requests.get(url) + """ + assert fires(code, "AI501") + + def test_aits11_chroma_retrieval_taint_flow(self): + code = """ + docs = vectorstore.similarity_search(user_query) + prompt = langchain.prompts.PromptTemplate.from_template(docs[0].page_content) + """ + assert fires(code, "AI101") + + def test_aits11_no_retrieval_no_taint(self): + code = """ + docs = [{"page_content": user_query}] + prompt = langchain.prompts.PromptTemplate.from_template(docs[0]["page_content"]) + """ + assert not fires(code, "AI101") def test_ai602_metadata(self): rule = _ai_rule("AI602") assert rule["severity"] == "Critical" assert rule["cwe"] == "CWE-78" - def test_ai602_pattern_matches(self): + def test_ai602_pattern_matches_shell_true(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_ai602_shell_false_is_safe(self): + code = """ + subprocess.run(["python", "-m", "pytest"]) + """ + assert not fires(code, "AI602") - def test_ai604_metadata(self): - rule = _ai_rule("AI604") - assert rule["severity"] == "High" - assert rule["cwe"] == "CWE-94" + def test_ai603_removed(self): + rule_ids = [rule["id"] for rule in toml.loads(RULES_PATH.read_text(encoding="utf-8"))["rule"]] + assert "AI603" not in rule_ids + + def test_ai604_removed(self): + rule_ids = [rule["id"] for rule in toml.loads(RULES_PATH.read_text(encoding="utf-8"))["rule"]] + assert "AI604" not in rule_ids def test_ai605_metadata(self): rule = _ai_rule("AI605") @@ -252,6 +296,12 @@ def test_ai605_pattern_matches(self): rule = _ai_rule("AI605") assert re.search(rule["pattern"], "agent = initialize_agent(verbose=True)") + def test_ai605_verbose_false_is_safe(self): + code = """ + agent = initialize_agent(verbose=False) + """ + assert not fires(code, "AI605") + def test_ai606_metadata(self): rule = _ai_rule("AI606") assert rule["severity"] == "Medium" @@ -276,16 +326,42 @@ def test_ai701_pattern_matches(self): rule = _ai_rule("AI701") assert re.search(rule["pattern"], "loader = DirectoryLoader('./docs')") + def test_ai701_plain_file_read_is_safe(self): + code = """ + content = Path("./data.txt").read_text() + """ + assert not fires(code, "AI701") + def test_ai702_metadata(self): rule = _ai_rule("AI702") assert rule["severity"] == "Medium" assert rule["cwe"] == "CWE-20" + def test_ai702_pattern_matches_score_threshold(self): + rule = _ai_rule("AI702") + assert re.search(rule["pattern"], "docs = vectorstore.similarity_search(query, score_threshold=0.3)") + + def test_ai702_top_k_without_threshold_is_safe(self): + code = """ + docs = vectorstore.similarity_search(query, k=4) + """ + assert not fires(code, "AI702") + def test_ai703_metadata(self): rule = _ai_rule("AI703") assert rule["severity"] == "Medium" assert rule["cwe"] == "CWE-400" + def test_ai703_pattern_matches(self): + rule = _ai_rule("AI703") + assert re.search(rule["pattern"], "chain = combine_docs(llm, docs)") + + def test_ai703_plain_context_join_is_safe(self): + code = """ + context = "\\n".join(docs) + """ + assert not fires(code, "AI703") + def test_ai704_metadata(self): rule = _ai_rule("AI704") assert rule["severity"] == "High" @@ -295,6 +371,12 @@ def test_ai704_pattern_matches(self): rule = _ai_rule("AI704") assert re.search(rule["pattern"], "embeddings = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')") + def test_ai704_other_embedding_model_is_safe(self): + code = """ + embeddings = OpenAIEmbeddings(model="text-embedding-3-small") + """ + assert not fires(code, "AI704") + # ------------------------------------------- # Tests for AI800 - API Key Management @@ -355,20 +437,68 @@ def test_ai901_pattern_matches(self): 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 not re.search(rule["pattern"], code) assert re.search(rule["exclude_pattern"], code) + def test_ai901_pattern_requires_plain_load(self): + rule = _ai_rule("AI901") + assert re.search(rule["pattern"], "data = yaml.load(llm_output, Loader=yaml.FullLoader)") + def test_ai902_metadata(self): rule = _ai_rule("AI902") assert rule["severity"] == "Medium" assert rule["cwe"] == "CWE-400" + def test_ai902_pattern_matches_llm_output_variable(self): + rule = _ai_rule("AI902") + assert re.search(rule["pattern"], "data = json.loads(output)") + + def test_ai902_plain_json_parse_is_safe(self): + code = """ + data = json.loads(request_body) + """ + assert not fires(code, "AI902") + def test_ai903_metadata(self): rule = _ai_rule("AI903") assert rule["severity"] == "Critical" assert rule["cwe"] == "CWE-94" + def test_ai903_pattern_matches_eval_of_llm_output(self): + rule = _ai_rule("AI903") + assert re.search(rule["pattern"], "result = eval(llm_output)") + + def test_ai903_plain_eval_is_safe(self): + code = """ + result = eval(expr) + """ + assert not fires(code, "AI903") + + def test_ai903_response_variable_line_is_safe(self): + code = """ + response = fetch_data() + print(response) + """ + assert not fires(code, "AI903") + def test_ai904_metadata(self): rule = _ai_rule("AI904") assert rule["severity"] == "High" assert rule["cwe"] == "CWE-79" + + def test_ai904_pattern_matches_inner_html(self): + rule = _ai_rule("AI904") + assert re.search(rule["pattern"], "element.innerHTML = llm_reply") + assert not re.search(rule["exclude_pattern"], "element.innerHTML = llm_reply") + + def test_ai904_sanitized_inner_html_is_safe(self): + code = """ + element.innerHTML = sanitize_html(llm_reply) + """ + assert not fires(code, "AI904") + + def test_ai904_response_variable_line_is_safe(self): + code = """ + response = render_template("index.html", data=data) + """ + assert not fires(code, "AI904")