"""
Check for insecure deserialization vulnerabilities.
"""

import os
import re
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple, Set

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


class InsecureDeserializationCheck(BaseCheck):
    """
    Check for insecure deserialization vulnerabilities.
    """
    
    name = "insecure-deserialization"
    description = "Check for insecure deserialization vulnerabilities"
    
    # Insecure deserialization patterns
    INSECURE_DESERIALIZATION_PATTERNS = [
        # pickle
        (r"pickle\.loads?\s*\(.*?\)", "Python pickle", "critical", "CWE-502", 
         "The pickle.loads function can execute arbitrary code when deserializing objects. "
         "Use safer_pickle or JSON for processing untrusted data."),
        
        # yaml.load without safe=True
        (r"yaml\.load\s*\((?!.*safe=True).*?\)", "YAML unsafe load", "high", "CWE-502",
         "The yaml.load method without the safe=True parameter can execute arbitrary code. "
         "Use yaml.safe_load instead."),
        
        # Marshal
        (r"marshal\.loads?\s*\(.*?\)", "Python marshal", "high", "CWE-502",
         "The marshal module is unsafe for untrusted data as it can execute arbitrary code."),
        
        # shelve (uses pickle)
        (r"shelve\.open\s*\(.*?\)", "Python shelve", "medium", "CWE-502",
         "The shelve module uses pickle and is unsafe for untrusted data."),
        
        # jsonpickle
        (r"jsonpickle\.(?:decode|loads)\s*\(.*?\)", "jsonpickle decode", "medium", "CWE-502",
         "jsonpickle can execute arbitrary code when deserializing objects."),
        
        # dill
        (r"dill\.loads?\s*\(.*?\)", "Python dill", "high", "CWE-502",
         "The dill module is unsafe and can execute arbitrary code during deserialization."),
        
        # pyyaml with constructors
        (r"yaml\..*?[Ll]oader.*?\)", "YAML with loader", "high", "CWE-502",
         "Loader in PyYAML can execute arbitrary code. Use yaml.SafeLoader."),
        
        # subprocess with shell=True
        (r"subprocess\.(?:call|run|Popen)\s*\(.*?shell\s*=\s*True.*?\)", "subprocess with shell=True", "high", "CWE-78",
         "Using subprocess with shell=True and input data can lead to command injection."),
    ]
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__(config)
    
    def run(self, target_path: str, **kwargs) -> Tuple[CheckResult, List[VulnerabilityIssue]]:
        """
        Run check for insecure deserialization.
        
        Args:
            target_path: Path to directory or file for checking
            **kwargs: Additional parameters
            
        Returns:
            Tuple[CheckResult, List[VulnerabilityIssue]]: Check result and list of found vulnerabilities
        """
        self.issues = []
        
        try:
            # Check files for insecure deserialization
            path = Path(target_path)
            for file_path in self._find_python_files(path):
                self._check_file_for_insecure_deserialization(file_path)
            
            # Create result
            status = CheckStatus.PASSED if not self.issues else CheckStatus.FAILED
            result = self.create_result(
                status=status,
                message=f"Found {len(self.issues)} potential insecure deserialization vulnerabilities" if self.issues else "No insecure deserialization vulnerabilities detected",
            )
            
            return result, self.issues
            
        except Exception as e:
            # In case of error, return ERROR status
            result = self.create_result(
                status=CheckStatus.ERROR,
                message=f"Error checking for insecure deserialization: {str(e)}",
            )
            return result, self.issues
    
    def _find_python_files(self, path: Path) -> Set[Path]:
        """
        Find Python files for checking.
        
        Args:
            path: Path to directory or file
            
        Returns:
            Set[Path]: Set of paths to files
        """
        result = set()
        
        if path.is_file() and path.suffix == ".py":
            # Check for exclusions
            if not self.should_exclude_path(path):
                result.add(path)
            return result
        
        # Traverse directory
        for root, dirs, files in os.walk(path):
            # Filter directories to exclude
            dirs[:] = [d for d in dirs if not self.should_exclude_path(os.path.join(root, d))]
            
            for file in files:
                if file.endswith(".py"):
                    file_path = Path(root) / file
                    if not self.should_exclude_path(file_path):
                        result.add(file_path)
        
        return result
    
    def _check_file_for_insecure_deserialization(self, file_path: Path) -> None:
        """
        Check file for insecure deserialization vulnerabilities.
        
        Args:
            file_path: Path to file
        """
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                content = f.read()
            
            lines = content.splitlines()
            
            # Compile regular expressions
            for i, line in enumerate(lines):
                for pattern_tuple in self.INSECURE_DESERIALIZATION_PATTERNS:
                    pattern, issue_type, severity, cwe_id, remediation = pattern_tuple
                    matches = re.finditer(pattern, line)
                    for match in matches:
                        # Check context to reduce false positives
                        if self._is_likely_vulnerable_context(lines, i, match.group(0)):
                            # Add vulnerability to the list
                            self.add_issue(
                                issue_type=f"Insecure Deserialization: {issue_type}",
                                severity=severity,
                                description=f"Potential insecure deserialization vulnerability: {issue_type}",
                                file_path=str(file_path),
                                line_number=i + 1,
                                code_snippet=line.strip(),
                                confidence="MEDIUM",
                                cwe_id=cwe_id,
                                owasp_category="A8:2017-Insecure Deserialization",
                                remediation=remediation,
                            )
        
        except Exception as e:
            # Skip files with errors
            print(f"Error checking file {file_path}: {str(e)}")
    
    def _is_likely_vulnerable_context(self, lines: List[str], line_index: int, match_string: str) -> bool:
        """
        Check context to determine probability of a real vulnerability.
        
        Args:
            lines: List of file lines
            line_index: Current line index
            match_string: Found match
            
        Returns:
            bool: True if context indicates a probable vulnerability
        """
        # Get several lines before and after for context analysis
        start = max(0, line_index - 3)
        end = min(len(lines), line_index + 4)
        context = "\n".join(lines[start:end])
        
        # Safe context indicators
        safe_indicators = [
            "# nosec",
            "# noqa",
            "def test_",
            "class Test",
            "assert",
            "unittest",
            "pytest",
        ]
        
        # Check for safe context indicators
        for indicator in safe_indicators:
            if indicator in context:
                return False
        
        # Vulnerable context indicators (data from untrusted sources)
        vulnerable_indicators = [
            "request", "user", "input", "data", "file", "http", "response", "get", "post", "json"
        ]
        
        # If it's pickle, yaml without safe=True, eval - they are always potentially dangerous
        if any(x in match_string.lower() for x in ["pickle", "yaml.load", "eval", "marshal"]):
            # Check for untrusted data sources nearby
            return any(indicator in context.lower() for indicator in vulnerable_indicators)
        
        return True 