"""
Module for detecting code injection vulnerabilities.
"""

import ast
import logging
import os
import re
from typing import Dict, List, Tuple, Set, Any

from vulnerability_checker.checks.base import BaseCheck
from vulnerability_checker.models import CheckResult, CheckStatus, VulnerabilityIssue, Severity


class CodeInjectionCheck(BaseCheck):
    """
    Check for code injection vulnerabilities, specifically looking for
    unsafe use of eval(), exec(), and other dangerous functions with
    user input.
    """

    def __init__(self, config: Dict[str, Any] = None):
        """
        Initialize the code injection check.
        
        Args:
            config: Check configuration
        """
        super().__init__(config)
        self.logger = logging.getLogger(__name__)
        self.dangerous_functions = {"eval", "exec", "compile", "__import__"}
        
        # Get custom patterns from config or use defaults
        if config and "additional_dangerous_patterns" in config:
            self.additional_patterns = config["additional_dangerous_patterns"]
        else:
            self.additional_patterns = [
                r"subprocess\.(?:Popen|call|check_output|check_call|run)\s*\(",
                r"os\.(?:system|popen|execl|execle|execlp|execlpe|execv|execve|execvp|execvpe)\s*\("
            ]
    
    def run(self, target_path: str, verbose: bool = False) -> Tuple[CheckResult, List[VulnerabilityIssue]]:
        """
        Run code injection vulnerability checks on the target path.
        
        Args:
            target_path: Path to the directory or file to check
            verbose: Enable verbose output
            
        Returns:
            Tuple[CheckResult, List[VulnerabilityIssue]]: Check result and list of found vulnerabilities
        """
        vulnerabilities = []
        
        if os.path.isfile(target_path) and target_path.endswith('.py'):
            self._check_file(target_path, vulnerabilities, verbose)
        elif os.path.isdir(target_path):
            for root, _, files in os.walk(target_path):
                for file in files:
                    if file.endswith('.py'):
                        file_path = os.path.join(root, file)
                        self._check_file(file_path, vulnerabilities, verbose)
        
        # Create check result
        status = CheckStatus.FAILED if vulnerabilities else CheckStatus.PASSED
        message = f"Found {len(vulnerabilities)} code injection vulnerabilities" if vulnerabilities else "No code injection vulnerabilities found"
        
        result = CheckResult(
            status=status,
            issues_count=len(vulnerabilities),
            message=message
        )
        
        return result, vulnerabilities
    
    def _check_file(self, file_path: str, vulnerabilities: List[VulnerabilityIssue], verbose: bool) -> None:
        """
        Check a single file for code injection vulnerabilities.
        
        Args:
            file_path: Path to the file to check
            vulnerabilities: List to append found vulnerabilities
            verbose: Enable verbose output
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # First, try to parse with AST for more accurate detection
            try:
                self._check_with_ast(file_path, content, vulnerabilities)
            except SyntaxError:
                # If AST parsing fails, fallback to regex
                self.logger.warning(f"AST parsing failed for {file_path}, falling back to regex")
                self._check_with_regex(file_path, content, vulnerabilities)
                
        except Exception as e:
            self.logger.error(f"Error checking file {file_path}: {str(e)}")
    
    def _check_with_ast(self, file_path: str, content: str, vulnerabilities: List[VulnerabilityIssue]) -> None:
        """
        Check for code injection using AST parsing.
        
        Args:
            file_path: Path to the file being checked
            content: File content
            vulnerabilities: List to append found vulnerabilities
        """
        tree = ast.parse(content)
        
        # Track variables that might contain user input
        user_input_vars = set()
        
        # First pass: identify variables that might contain user input
        for node in ast.walk(tree):
            # Look for assignments from request.data or similar
            if isinstance(node, ast.Assign):
                if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
                    var_name = node.targets[0].id
                    # Check if value comes from request data or similar
                    if isinstance(node.value, ast.Call):
                        if self._is_request_data_getter(node.value):
                            user_input_vars.add(var_name)
                    elif isinstance(node.value, ast.Attribute) and self._is_request_data(node.value):
                        user_input_vars.add(var_name)
        
        # Second pass: check for dangerous function calls
        for node in ast.walk(tree):
            if isinstance(node, ast.Call):
                # Check for direct calls to dangerous functions
                if isinstance(node.func, ast.Name) and node.func.id in self.dangerous_functions:
                    line_number = node.lineno
                    code_line = content.splitlines()[line_number - 1]
                    
                    # Look for indicators that user input is being used
                    if self._check_for_user_input(node, tree, user_input_vars):
                        self._add_vulnerability(file_path, line_number, code_line, node.func.id, vulnerabilities)
                
                # Check for attribute calls like "subprocess.call"
                elif isinstance(node.func, ast.Attribute):
                    if hasattr(node.func, 'attr'):
                        # Check for subprocess and os command execution
                        if (hasattr(node.func, 'value') and 
                            isinstance(node.func.value, ast.Name) and
                            node.func.value.id in {"subprocess", "os"} and
                            node.func.attr in {"system", "popen", "call", "check_output", "Popen", "run"}):
                            
                            line_number = node.lineno
                            code_line = content.splitlines()[line_number - 1]
                            
                            if self._check_for_user_input(node, tree, user_input_vars):
                                function_name = f"{node.func.value.id}.{node.func.attr}"
                                self._add_vulnerability(file_path, line_number, code_line, function_name, vulnerabilities)
    
    def _is_request_data_getter(self, node: ast.Call) -> bool:
        """
        Check if the call is a getter for request data like request.data.get()
        """
        if isinstance(node.func, ast.Attribute) and node.func.attr == "get":
            if isinstance(node.func.value, ast.Attribute):
                if node.func.value.attr in {"data", "form", "args", "json", "params", "cookies", "files"}:
                    if isinstance(node.func.value.value, ast.Name) and node.func.value.value.id == "request":
                        return True
        return False
    
    def _check_for_user_input(self, node: ast.Call, tree: ast.Module, user_input_vars: Set[str] = None) -> bool:
        """
        Check if the function call seems to be using user input.
        
        Args:
            node: The AST node representing the function call
            tree: The entire AST tree
            user_input_vars: Set of variable names identified as potentially containing user input
            
        Returns:
            bool: True if the call appears to use user input
        """
        if user_input_vars is None:
            user_input_vars = set()
            
        # Check if any argument looks like it's coming from user input
        for arg in node.args:
            # Check for request.* or input()
            if (isinstance(arg, ast.Call) and 
                isinstance(arg.func, ast.Name) and 
                arg.func.id == "input"):
                return True
            
            # Check for request.data.get calls
            if isinstance(arg, ast.Call) and self._is_request_data_getter(arg):
                return True
            
            # Check for request.data, request.form, etc.
            if isinstance(arg, ast.Attribute) and self._is_request_data(arg):
                return True
                
            # Check for variable references that might be user input
            if isinstance(arg, ast.Name):
                # Check against our identified user input variables
                if arg.id in user_input_vars:
                    return True
                    
                # Check if variable name suggests user input
                user_input_indicators = {"user", "input", "request", "param", "data", "form", "json", "post"}
                if any(indicator in arg.id.lower() for indicator in user_input_indicators):
                    return True
        
        # Check for keyword arguments too
        for keyword in node.keywords:
            if isinstance(keyword.value, ast.Attribute) and self._is_request_data(keyword.value):
                return True
            
            if isinstance(keyword.value, ast.Call) and self._is_request_data_getter(keyword.value):
                return True
            
            if isinstance(keyword.value, ast.Name) and keyword.value.id in user_input_vars:
                return True
        
        return False
    
    def _is_request_data(self, node: ast.Attribute) -> bool:
        """
        Check if an attribute access is related to request data.
        
        Args:
            node: The AST node
            
        Returns:
            bool: True if the node appears to access request data
        """
        if isinstance(node.value, ast.Attribute):
            if hasattr(node.value, 'attr') and node.value.attr in {"data", "form", "args", "json", "params", "cookies", "files"}:
                if isinstance(node.value.value, ast.Name) and node.value.value.id == "request":
                    return True
        elif isinstance(node.value, ast.Name) and node.value.id == "request":
            if node.attr in {"data", "form", "args", "json", "params", "cookies", "files"}:
                return True
        return False
    
    def _check_with_regex(self, file_path: str, content: str, vulnerabilities: List[VulnerabilityIssue]) -> None:
        """
        Check for code injection using regex patterns as a fallback.
        
        Args:
            file_path: Path to the file being checked
            content: File content
            vulnerabilities: List to append found vulnerabilities
        """
        lines = content.splitlines()
        
        # Check for dangerous function calls
        for i, line in enumerate(lines):
            # Check direct eval/exec calls with request data
            for func in self.dangerous_functions:
                # Pattern to match func(request...) or func(something_request_related)
                pattern = fr"{func}\s*\(.*(?:request|input|user|param|data|form|json|post|get).*\)"
                if re.search(pattern, line, re.IGNORECASE):
                    self._add_vulnerability(file_path, i + 1, line, func, vulnerabilities)
                
                # Specific pattern for request.data.get() cases
                data_get_pattern = fr"{func}\s*\(.+?request\.[\w]+\.get\(.+?\).+?\)"
                if re.search(data_get_pattern, line):
                    self._add_vulnerability(file_path, i + 1, line, func, vulnerabilities)
                
                # Look for variables that were assigned from request data
                # This is a simplified heuristic - in a real implementation, we'd need more context
                context_lines = lines[max(0, i-5):i]
                for context_line in context_lines:
                    # Look for variable assignments from request data
                    var_assign_match = re.search(r"(\w+)\s*=\s*request\.[\w]+\.get\(.+?\)", context_line)
                    if var_assign_match:
                        var_name = var_assign_match.group(1)
                        # Check if that variable is used in the current line with the dangerous function
                        if re.search(fr"{func}\s*\(\s*{var_name}", line):
                            self._add_vulnerability(file_path, i + 1, line, func, vulnerabilities)
            
            # Check additional patterns
            for pattern in self.additional_patterns:
                if re.search(pattern, line) and re.search(r"(?:request|input|user|param|data|form|json|post|get)", line, re.IGNORECASE):
                    self._add_vulnerability(file_path, i + 1, line, "command execution", vulnerabilities)
    
    def _add_vulnerability(self, file_path: str, line_number: int, code_line: str, function_name: str, vulnerabilities: List[VulnerabilityIssue]) -> None:
        """
        Add a code injection vulnerability to the list.
        
        Args:
            file_path: Path to the file
            line_number: Line number in the file
            code_line: The vulnerable line of code
            function_name: The dangerous function being used
            vulnerabilities: List to append the vulnerability
        """
        vulnerability = VulnerabilityIssue(
            issue_type=f"Code Injection: Unsafe use of {function_name}",
            severity=Severity.HIGH.value,
            description=(
                f"Found potentially unsafe use of {function_name} with user-controlled input. "
                f"This can lead to code injection attacks allowing attackers to execute arbitrary code."
            ),
            file_path=file_path,
            line_number=line_number,
            code_snippet=code_line.strip(),
            confidence="HIGH",
            cwe_id="CWE-94",
            owasp_category="A1:2017-Injection",
            remediation=(
                f"Avoid using {function_name} with user input. Instead, use safer alternatives such as:\n"
                "1. For dynamic evaluation, use ast.literal_eval() which only evaluates literals\n"
                "2. Implement a custom safe parser for more complex expressions\n"
                "3. Use specific functions for the intended purpose rather than generic execution"
            ),
            check_name="code-injection"
        )
        
        vulnerabilities.append(vulnerability) 