"""TS-MCP-004: Prompt Injection for TypeScript/JavaScript.

Detects patterns that could enable prompt injection attacks:
- User input concatenated into LLM prompts
- Dynamic tool descriptions
"""

from __future__ import annotations

import re

from . import TSRule
from ..models import Finding, Severity

# Tool registration with dynamic description
TOOL_DYNAMIC_DESC = re.compile(
    r'(?:\.tool|\.addTool|\.registerTool|server\.tool)\s*\(\s*'
    r'[^,]+,\s*(?:`[^`]*\$\{|[^,]*\+)'
)

# LLM API call patterns
LLM_CALL = re.compile(
    r'(?:\.create\s*\(|\.generate\s*\(|\.complete\s*\(|'
    r'\.chat\s*\(|\.invoke\s*\(|'
    r'openai\.|anthropic\.|'
    r'ChatCompletion|'
    r'generateText|streamText)'
)

# User input in prompt
PROMPT_CONCAT = re.compile(
    r'(?:prompt|message|content|system|instruction)\s*'
    r'(?:[:=]\s*`[^`]*\$\{|[:=]\s*[^;]*\+\s*(?:args|params|input|'
    r'req\.|body\.|query\.|user))'
)

# Template literal in prompt context
PROMPT_TEMPLATE = re.compile(
    r'(?:role|content|text)\s*:\s*`[^`]*\$\{[^}]*'
    r'(?:args|params|input|user|query|body)'
)


class TSPromptInjectionRule(TSRule):
    """Detects prompt injection vectors in TS/JS MCP servers."""

    @property
    def rule_id(self) -> str:
        return "TS-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 (
            "User input concatenated into LLM prompts or dynamic "
            "tool descriptions may enable prompt injection."
        )

    @property
    def detailed_description(self) -> str:
        return (
            "User input concatenated into LLM prompts or dynamic "
            "tool descriptions may enable prompt injection.\n\n"
            "Examples of vulnerable code:\n"
            "  const prompt = `Analyze: ${userInput}`;\n"
            "  server.tool('query', `Search for ${desc}`, ...);\n\n"
            "How to fix:\n"
            "  - Use static tool descriptions\n"
            "  - Separate user content from system instructions\n"
            "  - Use structured message formats\n"
            "  - Sanitize/escape user input in prompts"
        )

    def check(self, source: str) -> list[Finding]:
        """Check for prompt injection patterns."""
        findings: list[Finding] = []
        lines = self._get_lines(source)

        for i, line in enumerate(lines, 1):
            if self._is_comment_line(line):
                continue

            # Dynamic tool description
            if TOOL_DYNAMIC_DESC.search(line):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Dynamic Tool Description",
                    description=(
                        "Tool registered with dynamic description "
                        "(prompt injection risk)."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip()[:80],
                    recommendation=(
                        "Use static string literals for tool "
                        "descriptions."
                    ),
                    cwe=self.cwe,
                ))
                continue

            # User input in prompt template
            if PROMPT_TEMPLATE.search(line):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="User Input in LLM Prompt",
                    description=(
                        "User-controlled data interpolated into "
                        "LLM prompt content."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip()[:80],
                    recommendation=(
                        "Separate user content from system "
                        "instructions using structured messages."
                    ),
                    cwe=self.cwe,
                ))
                continue

            # Prompt concatenation with user input
            if PROMPT_CONCAT.search(line):
                findings.append(Finding(
                    rule_id=self.rule_id,
                    severity=self.severity,
                    title="Prompt Concatenation with User Input",
                    description=(
                        "User input concatenated into prompt string."
                    ),
                    file="",
                    line=i,
                    code_snippet=line.strip()[:80],
                    recommendation=(
                        "Use structured message arrays instead of "
                        "string concatenation."
                    ),
                    cwe=self.cwe,
                    confidence="MEDIUM",
                ))

        return findings
