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

Detects patterns that could enable prompt injection attacks in MCP servers:
- Dynamic tool descriptions (f-strings, format() in @tool decorators)
- User input concatenated directly into LLM prompts
"""

from __future__ import annotations

import ast

from . import Rule
from ..models import Finding, Severity

# Function names that send prompts to LLMs
LLM_CALL_FUNCS = {
    "complete", "chat", "generate", "invoke", "create",
    "messages.create", "chat.completions.create",
    "completions.create", "send_message",
}

# Attribute names on LLM client objects
LLM_ATTRS = {
    "complete", "chat", "generate", "invoke", "create",
    "send_message", "run",
}


class PromptInjectionRule(Rule):
    """Detects prompt injection vectors in MCP tool definitions.

    Flags two patterns:
    1. Dynamic tool descriptions (f-strings or .format() in decorators)
    2. User input concatenated directly into LLM prompts
    """

    @property
    def rule_id(self) -> str:
        return "MCP-004"

    @property
    def name(self) -> str:
        return "Prompt Injection Vector"

    @property
    def severity(self) -> Severity:
        return Severity.HIGH

    @property
    def cwe(self) -> str:
        return "CWE-74"

    @property
    def description(self) -> str:
        return (
            "Tool description or prompt construction may enable "
            "prompt injection attacks."
        )

    @property
    def detailed_description(self) -> str:
        return (
            "Tool description or prompt construction may enable "
            "prompt injection attacks.\n\n"
            "Prompt injection occurs when user-controlled data is embedded\n"
            "into LLM prompts without sanitization. In MCP, this can happen\n"
            "through dynamic tool descriptions or through tool handlers that\n"
            "concatenate user input into prompts.\n\n"
            "How to fix:\n"
            "  - Use static string literals for tool descriptions\n"
            "  - Separate user content from system instructions\n"
            "  - Use structured message formats instead of string concat\n"
            "  - Sanitize user input before embedding in prompts"
        )

    def check(self, tree: ast.Module, source: str) -> list[Finding]:
        """Check for prompt injection patterns."""
        findings: list[Finding] = []
        findings.extend(self._check_dynamic_descriptions(tree, source))
        findings.extend(self._check_prompt_concatenation(tree, source))
        return findings

    def _check_dynamic_descriptions(
        self, tree: ast.Module, source: str
    ) -> list[Finding]:
        """Detect f-strings or format() in tool decorator arguments."""
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            for dec in node.decorator_list:
                if not self._is_tool_decorator(dec):
                    continue
                dynamic_nodes = self._find_dynamic_strings(dec)
                for dyn_node in dynamic_nodes:
                    findings.append(Finding(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        title="Dynamic tool description (prompt injection risk)",
                        description=(
                            f"Tool decorator for '{node.name}' uses a dynamic "
                            f"string in its description. An attacker could "
                            f"manipulate the tool's behavior via injected content."
                        ),
                        file="",
                        line=getattr(dyn_node, "lineno", node.lineno),
                        code_snippet=self._get_snippet(
                            source, getattr(dyn_node, "lineno", node.lineno)
                        ),
                        recommendation=(
                            "Use static string literals for tool descriptions."
                        ),
                        cwe=self.cwe,
                        confidence="HIGH",
                        remediation=(
                            "Replace the dynamic string (f-string or .format()) "
                            "in the tool decorator with a static string literal. "
                            "Tool descriptions are consumed by the LLM and must "
                            "not contain user-controlled content."
                        ),
                    ))
        return findings

    def _check_prompt_concatenation(
        self, tree: ast.Module, source: str
    ) -> list[Finding]:
        """Detect user input concatenated into LLM prompts."""
        findings: list[Finding] = []
        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not self._is_tool_handler(node):
                continue

            param_names = {arg.arg for arg in node.args.args}
            for child in ast.walk(node):
                if not isinstance(child, ast.Call):
                    continue
                if not self._is_llm_call(child):
                    continue
                if self._call_uses_dynamic_input(child, param_names):
                    findings.append(Finding(
                        rule_id=self.rule_id,
                        severity=self.severity,
                        title="User input in LLM prompt (prompt injection)",
                        description=(
                            f"Tool handler '{node.name}' concatenates user "
                            f"input into an LLM prompt call."
                        ),
                        file="",
                        line=child.lineno,
                        code_snippet=self._get_snippet(source, child.lineno),
                        recommendation=(
                            "Separate user content from system instructions. "
                            "Use structured message formats."
                        ),
                        cwe=self.cwe,
                        confidence="MEDIUM",
                        remediation=(
                            "Instead of concatenating user input into the "
                            "prompt string, pass it as a separate 'user' "
                            "message in a structured messages array. This "
                            "prevents user content from overriding system "
                            "instructions."
                        ),
                    ))
        return findings

    def _is_tool_decorator(self, dec: ast.expr) -> bool:
        """Check if a decorator is a tool registration decorator."""
        dec_str = ast.dump(dec).lower()
        return "tool" in dec_str

    def _find_dynamic_strings(self, node: ast.AST) -> list[ast.AST]:
        """Find f-strings or .format() calls within an AST subtree."""
        results: list[ast.AST] = []
        for child in ast.walk(node):
            if isinstance(child, ast.JoinedStr):
                results.append(child)
            if isinstance(child, ast.Call):
                if (
                    isinstance(child.func, ast.Attribute)
                    and child.func.attr == "format"
                ):
                    results.append(child)
        return results

    def _is_tool_handler(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> bool:
        """Check if function is an MCP tool handler."""
        name_lower = node.name.lower()
        if any(kw in name_lower for kw in ("tool", "handle")):
            return True
        for dec in node.decorator_list:
            if "tool" in ast.dump(dec).lower():
                return True
        return False

    def _is_llm_call(self, node: ast.Call) -> bool:
        """Check if a call node looks like an LLM API call."""
        func_name = self._get_func_name(node)
        if func_name in LLM_CALL_FUNCS:
            return True
        # Check for client.method patterns
        if isinstance(node.func, ast.Attribute):
            if node.func.attr in LLM_ATTRS:
                return True
        return False

    def _call_uses_dynamic_input(
        self, call: ast.Call, param_names: set[str]
    ) -> bool:
        """Check if an LLM call uses dynamic strings with param references."""
        for arg in call.args:
            if self._has_dynamic_param_ref(arg, param_names):
                return True
        for kw in call.keywords:
            if kw.value and self._has_dynamic_param_ref(
                kw.value, param_names
            ):
                return True
        return False

    def _has_dynamic_param_ref(
        self, node: ast.AST, param_names: set[str]
    ) -> bool:
        """Check if node is a dynamic string referencing function params."""
        if isinstance(node, ast.JoinedStr):
            for child in ast.walk(node):
                if isinstance(child, ast.Name) and child.id in param_names:
                    return True
        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
            for child in ast.walk(node):
                if isinstance(child, ast.Name) and child.id in param_names:
                    return True
        return False
