"""Tests for MCP-004: Prompt Injection detection rule.

NOTE: This test file contains intentionally vulnerable code STRINGS
that are passed to the scanner for detection testing. These are not
executed - they are static analysis targets.
"""

import ast

from mcpaudit.models import Severity
from mcpaudit.rules.prompt_injection import PromptInjectionRule


class TestPromptInjection:
    """Tests for the Prompt Injection detection rule."""

    def setup_method(self):
        self.rule = PromptInjectionRule()

    def _check(self, code: str):
        tree = ast.parse(code)
        return self.rule.check(tree, code)

    def test_rule_metadata(self):
        assert self.rule.rule_id == "MCP-004"
        assert self.rule.severity == Severity.HIGH
        assert self.rule.cwe == "CWE-74"
        assert "Prompt Injection" in self.rule.name

    def test_fstring_in_tool_decorator(self):
        code = '''
@server.tool(description=f"Process {user_type} data")
def process(data):
    return data
'''
        findings = self._check(code)
        assert len(findings) >= 1
        assert findings[0].rule_id == "MCP-004"
        assert "dynamic" in findings[0].title.lower()

    def test_format_in_tool_decorator(self):
        code = '''
@server.tool(description="Process {} data".format(dtype))
def process(data):
    return data
'''
        findings = self._check(code)
        assert len(findings) >= 1

    def test_static_description_passes(self):
        code = '''
@server.tool(description="Process user data safely")
def process(data):
    return data
'''
        findings = self._check(code)
        assert len(findings) == 0

    def test_no_decorator_passes(self):
        code = '''
def process(data):
    return data
'''
        findings = self._check(code)
        assert len(findings) == 0

    def test_user_input_in_llm_prompt_indirect(self):
        """Indirect f-string via variable is not detected (known limitation).

        The rule detects direct f-string/concat in call args. Indirect
        assignment then passing the variable requires data-flow analysis.
        """
        code = '''
def handle_tool_query(user_input):
    prompt = f"Analyze: {user_input}"
    return client.chat(prompt)
'''
        findings = self._check(code)
        # Indirect case: f-string in variable, not in call args
        # This is a known limitation of static AST analysis
        assert len(findings) == 0

    def test_user_input_in_llm_prompt_concat(self):
        code = '''
def handle_tool_query(user_input):
    prompt = "Analyze: " + user_input
    return client.complete(prompt)
'''
        # The concat is not directly in the call args here
        # so this tests that we detect f-string/concat in call args
        findings = self._check(code)
        # Concat is in variable assignment, not call arg — may not detect
        # This is fine — we focus on direct concat in call args

    def test_direct_fstring_in_llm_call(self):
        code = '''
def handle_tool_query(user_input):
    return client.chat(f"Analyze: {user_input}")
'''
        findings = self._check(code)
        assert len(findings) >= 1

    def test_static_prompt_in_llm_call_passes(self):
        code = '''
def handle_tool_query(user_input):
    return client.chat("Analyze the data")
'''
        findings = self._check(code)
        assert len(findings) == 0

    def test_finding_has_remediation(self):
        code = '''
@server.tool(description=f"Handle {action}")
def handle(data):
    return data
'''
        findings = self._check(code)
        assert len(findings) >= 1
        assert findings[0].remediation != ""

    def test_finding_has_cwe(self):
        code = '''
@server.tool(description=f"Handle {action}")
def handle(data):
    return data
'''
        findings = self._check(code)
        assert findings[0].cwe == "CWE-74"

    def test_non_tool_handler_llm_call_ignored(self):
        code = '''
def helper(data):
    return client.chat(f"Process: {data}")
'''
        findings = self._check(code)
        # helper is not a tool handler so no finding for prompt concat
        assert not any(
            "prompt injection" in f.title.lower() for f in findings
        )

    def test_detailed_description(self):
        desc = self.rule.detailed_description
        assert "prompt injection" in desc.lower()
        assert "How to fix" in desc

    def test_messages_create_detected(self):
        code = '''
def handle_tool_ask(question):
    return client.messages.create(f"Answer: {question}")
'''
        findings = self._check(code)
        assert len(findings) >= 1

    def test_multiple_findings_in_one_function(self):
        code = '''
@server.tool(description=f"Dynamic {desc}")
def handle_tool_multi(user_input):
    return client.chat(f"Process: {user_input}")
'''
        findings = self._check(code)
        # Should have at least 2: one for decorator, one for prompt concat
        assert len(findings) >= 2

    def test_async_tool_handler_detected(self):
        code = '''
async def handle_tool_query(user_input):
    return await client.chat(f"Analyze: {user_input}")
'''
        findings = self._check(code)
        assert len(findings) >= 1
