"""
Injection Pattern Analyzer for MCP Security Auditor.

Detects prompt injection vulnerabilities and unsafe input handling patterns.
"""

import re
from pathlib import Path
from typing import List, Dict, Any

from ..core.config import AuditConfig
from ..core.models import (
    Finding, CodeLocation, Remediation, MCPServerInfo
)
from .static import BaseAnalyzer


class InjectionPatternAnalyzer(BaseAnalyzer):
    """
    Analyzes code for prompt injection vulnerabilities.
    
    Checks for:
    - Unvalidated user input in prompts
    - Missing input sanitization
    - Dangerous string concatenation patterns
    - Template injection risks
    - Command injection through prompts
    """
    
    name = "injection"
    
    # Patterns indicating prompt injection vulnerabilities
    INJECTION_PATTERNS = {
        "python": [
            # Direct concatenation with user input
            (r'f["\'].*\{.*input.*\}.*["\']', "high", 
             "User input directly interpolated in f-string"),
            (r'\.format\s*\([^)]*input[^)]*\)', "high",
             "User input in .format() call"),
            (r'%\s*\([^)]*input[^)]*\)', "medium",
             "User input in % string formatting"),
            
            # Prompt construction patterns
            (r'prompt\s*[+=]\s*.*user|user.*\+.*prompt', "high",
             "User input concatenated with prompt"),
            (r'messages\s*\.\s*append\s*\([^)]*user_input', "high",
             "User input directly added to messages"),
            (r'system_prompt\s*=.*\+.*input', "critical",
             "User input in system prompt construction"),
            
            # Template-based injection
            (r'Template\s*\([^)]*\)\.substitute\s*\([^)]*user', "high",
             "User input in Template.substitute()"),
            (r'jinja.*render.*user|user.*jinja.*render', "high",
             "User input in Jinja template rendering"),
            
            # Eval/exec with user input
            (r'eval\s*\([^)]*input', "critical",
             "User input in eval()"),
            (r'exec\s*\([^)]*input', "critical",
             "User input in exec()"),
            
            # XML/HTML injection
            (r'xml\..*\([^)]*user|user.*xml\.', "medium",
             "User input in XML construction"),
            (r'html\s*[+=].*user|user.*html', "medium",
             "User input in HTML construction"),
        ],
        "typescript": [
            # Template literal injection
            (r'`[^`]*\$\{[^}]*input[^}]*\}[^`]*`', "high",
             "User input in template literal"),
            (r'`[^`]*\$\{[^}]*user[^}]*\}[^`]*`', "high",
             "User variable in template literal"),
            
            # String concatenation
            (r'prompt\s*\+\s*.*input|input.*\+.*prompt', "high",
             "User input concatenated with prompt"),
            (r'messages\.push\s*\([^)]*user', "high",
             "User input directly pushed to messages"),
            
            # Dangerous patterns
            (r'eval\s*\([^)]*input', "critical",
             "User input in eval()"),
            (r'new\s+Function\s*\([^)]*input', "critical",
             "User input in Function constructor"),
            
            # DOM injection
            (r'innerHTML\s*=.*user|user.*innerHTML', "high",
             "User input in innerHTML assignment"),
            (r'document\.write\s*\([^)]*user', "high",
             "User input in document.write()"),
        ]
    }
    
    # Patterns that indicate safe input handling
    SAFE_PATTERNS = [
        r'sanitize',
        r'escape',
        r'validate',
        r'clean',
        r'strip_tags',
        r'html\.escape',
        r'markupsafe',
        r'bleach',
        r'DOMPurify',
    ]
    
    # MCP-specific injection risks
    MCP_INJECTION_PATTERNS = [
        # Tool input directly used in responses
        (r'return\s+.*params\.|return\s+.*args\.', "medium",
         "Tool parameters directly returned without processing"),
        
        # User content in tool descriptions
        (r'description\s*=.*\+.*user|user.*\+.*description', "high",
         "User input in tool description"),
        
        # Dynamic tool registration
        (r'register.*tool.*user|user.*register.*tool', "critical",
         "User input in dynamic tool registration"),
        
        # Resource URI injection
        (r'resource.*uri.*user|user.*resource.*uri', "high",
         "User input in resource URI"),
    ]
    
    async def analyze(self, path: str, server_info: MCPServerInfo) -> List[Finding]:
        """Analyze code for injection vulnerabilities."""
        findings = []
        path = Path(path)
        
        # Get language-specific patterns
        lang_patterns = self.INJECTION_PATTERNS.get(
            server_info.language,
            self.INJECTION_PATTERNS["python"]
        )
        
        # Determine file extensions
        if server_info.language == "python":
            extensions = ["*.py"]
        else:
            extensions = ["*.ts", "*.js", "*.mjs"]
        
        for ext in extensions:
            for file_path in path.rglob(ext):
                if self._should_skip_file(file_path):
                    continue
                
                content = self._read_file_content(file_path)
                if not content:
                    continue
                
                relative_path = str(file_path.relative_to(path))
                
                # Check language-specific patterns
                findings.extend(self._check_patterns(
                    file_path, content, relative_path, lang_patterns
                ))
                
                # Check MCP-specific patterns
                findings.extend(self._check_patterns(
                    file_path, content, relative_path, self.MCP_INJECTION_PATTERNS
                ))
                
                # Analyze input flow
                findings.extend(self._analyze_input_flow(
                    file_path, content, relative_path, server_info
                ))
        
        return findings
    
    def _check_patterns(
        self, 
        file_path: Path, 
        content: str, 
        relative_path: str,
        patterns: List[tuple]
    ) -> List[Finding]:
        """Check content against a list of patterns."""
        findings = []
        lines = content.split('\n')
        
        for pattern, severity, description in patterns:
            try:
                regex = re.compile(pattern, re.IGNORECASE)
                
                for i, line in enumerate(lines, 1):
                    if regex.search(line):
                        # Check if there's sanitization nearby
                        context_start = max(0, i - 5)
                        context_end = min(len(lines), i + 5)
                        context = '\n'.join(lines[context_start:context_end])
                        
                        has_sanitization = any(
                            re.search(safe, context, re.IGNORECASE) 
                            for safe in self.SAFE_PATTERNS
                        )
                        
                        # Reduce severity if sanitization is present
                        if has_sanitization:
                            if severity == "critical":
                                severity = "high"
                            elif severity == "high":
                                severity = "medium"
                            else:
                                continue  # Skip low severity with sanitization
                        
                        findings.append(Finding(
                            title=f"Potential injection vulnerability",
                            description=f"{description}. This could allow prompt injection attacks.",
                            severity=severity,
                            category="injection",
                            analyzer=self.name,
                            location=CodeLocation(
                                file=relative_path,
                                line_start=i,
                                line_end=i,
                                snippet=self._get_code_snippet(file_path, i)
                            ),
                            remediation=Remediation(
                                description="Validate and sanitize all user inputs before using them "
                                           "in prompts or string construction. Consider using "
                                           "structured input schemas and output parsers.",
                                code_example=self._get_remediation_example(description),
                                references=[
                                    "https://owasp.org/www-project-top-10-for-large-language-model-applications/",
                                    "https://cwe.mitre.org/data/definitions/94.html"
                                ]
                            ),
                            cwe_id="CWE-94",
                            confidence="high" if not has_sanitization else "medium"
                        ))
            except re.error:
                continue
        
        return findings
    
    def _analyze_input_flow(
        self,
        file_path: Path,
        content: str,
        relative_path: str,
        server_info: MCPServerInfo
    ) -> List[Finding]:
        """Analyze the flow of user input through the code."""
        findings = []
        lines = content.split('\n')
        
        # Find tool function definitions and their parameters
        if server_info.language == "python":
            tool_pattern = re.compile(r'@(?:mcp|server)\.tool')
            func_pattern = re.compile(r'async\s+def\s+(\w+)\s*\(([^)]*)\)')
        else:
            tool_pattern = re.compile(r'\.tool\s*\(|registerTool')
            func_pattern = re.compile(r'(?:async\s+)?(?:function\s+)?(\w+)\s*\(([^)]*)\)')
        
        in_tool_function = False
        current_tool_line = 0
        param_names = set()
        
        for i, line in enumerate(lines, 1):
            # Detect tool decorator/registration
            if tool_pattern.search(line):
                in_tool_function = True
                current_tool_line = i
                continue
            
            # Detect function definition after tool decorator
            if in_tool_function:
                func_match = func_pattern.search(line)
                if func_match:
                    # Extract parameter names
                    params = func_match.group(2)
                    param_names = set(re.findall(r'(\w+)\s*[,:\)]', params))
                    in_tool_function = False
                    
                    # Check the function body for unsafe parameter usage
                    j = i + 1
                    brace_count = 1 if '{' in line else 0
                    indent_level = len(line) - len(line.lstrip())
                    
                    while j < len(lines) and j < i + 100:  # Limit search
                        func_line = lines[j - 1]
                        
                        # Check for function end (Python: dedent, JS: closing brace)
                        if server_info.language == "python":
                            if func_line.strip() and not func_line.startswith(' ' * (indent_level + 1)):
                                if not func_line.strip().startswith('#'):
                                    break
                        else:
                            brace_count += func_line.count('{') - func_line.count('}')
                            if brace_count <= 0:
                                break
                        
                        # Check for dangerous patterns with parameters
                        for param in param_names:
                            if param in ['self', 'ctx', 'context', 'params', 'args']:
                                continue
                            
                            # Direct return of user input
                            if re.search(rf'return\s+.*{param}[^a-zA-Z_]', func_line):
                                # Check if it's being processed
                                if not any(safe in func_line.lower() for safe in self.SAFE_PATTERNS):
                                    findings.append(Finding(
                                        title=f"Tool returns unsanitized user input",
                                        description=f"Tool function directly returns parameter '{param}' "
                                                   f"without validation or sanitization.",
                                        severity="medium",
                                        category="injection",
                                        analyzer=self.name,
                                        location=CodeLocation(
                                            file=relative_path,
                                            line_start=j,
                                            line_end=j,
                                            snippet=self._get_code_snippet(file_path, j)
                                        ),
                                        remediation=Remediation(
                                            description="Validate and sanitize user input before returning. "
                                                       "Consider using Pydantic models for input validation."
                                        ),
                                        confidence="medium"
                                    ))
                        
                        j += 1
        
        return findings
    
    def _get_remediation_example(self, description: str) -> str:
        """Get a code example for remediation based on the issue."""
        description_lower = description.lower()
        
        if 'f-string' in description_lower or 'format' in description_lower:
            return """# Instead of:
prompt = f"Process this: {user_input}"

# Use validated input:
from pydantic import BaseModel, Field, field_validator

class UserInput(BaseModel):
    content: str = Field(..., max_length=1000)
    
    @field_validator('content')
    @classmethod
    def sanitize_content(cls, v):
        # Remove potential injection patterns
        v = v.replace('{', '').replace('}', '')
        return v.strip()

validated = UserInput(content=user_input)
prompt = f"Process this: {validated.content}"
"""
        
        elif 'system prompt' in description_lower:
            return """# NEVER include user input in system prompts
# Instead, use structured message roles:

messages = [
    {"role": "system", "content": FIXED_SYSTEM_PROMPT},
    {"role": "user", "content": sanitize(user_input)}
]
"""
        
        elif 'eval' in description_lower:
            return """# Never use eval() with user input
# Instead, use safe alternatives:

import ast
# For literal values only:
result = ast.literal_eval(user_input)

# Or use a proper parser for expressions
"""
        
        elif 'template' in description_lower:
            return """# Use safe templating with autoescape:
from jinja2 import Environment, select_autoescape

env = Environment(autoescape=select_autoescape(['html', 'xml']))
template = env.from_string(template_str)
result = template.render(user_input=escape(user_input))
"""
        
        return """# Validate and sanitize all user inputs:
def sanitize_input(value: str) -> str:
    # Remove potentially dangerous characters
    sanitized = re.sub(r'[<>{}\\[\\]|`$]', '', value)
    # Limit length
    return sanitized[:1000].strip()
"""
