import requests
import json
import time
from typing import List, Dict, Any, Optional, Union
import logging
from dosuby.src.interfaces.vulnerability_checker import VulnerabilityChecker

class NVDVulnerabilityChecker(VulnerabilityChecker):
    """
    A vulnerability checker implementation using the NVD (National Vulnerability Database) 
    JSON API.
    """
    
    def __init__(self, api_key: Optional[str] = None):
        """Initialize the NVD vulnerability checker.
        
        Args:
            api_key (str, optional): NVD API key for higher rate limits. Defaults to None.
        """
        self.api_base_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
        self.api_key = api_key
        self.logger = logging.getLogger("dosuby.vulnerability_checker")
        self.max_retries = 3
        self.retry_delay = 2  # seconds
        
        # Create a dict to cache results and avoid repeated requests
        self.cache = {}
        
    def check_cms_vulnerabilities(self, cms_name: str, version: str) -> List[Dict[str, Any]]:
        """Check vulnerabilities for a CMS version.
        
        Args:
            cms_name (str): Name of the CMS (e.g., "WordPress", "Drupal").
            version (str): Version number of the CMS.
            
        Returns:
            List[Dict[str, Any]]: List of vulnerability information.
        """
        if not cms_name or not version:
            return []
            
        # Create a cache key for this query
        cache_key = f"{cms_name.lower()}-{version}"
        
        # Check cache first
        if cache_key in self.cache:
            self.logger.info(f"Using cached vulnerability data for {cms_name} {version}")
            return self.cache[cache_key]
            
        # Normalize CMS name for searching
        cms_keywords = {
            "wordpress": ["wordpress", "wp"],
            "drupal": ["drupal"],
            "joomla": ["joomla"],
            "moodle": ["moodle"],
            "magento": ["magento"],
            "typo3": ["typo3"],
            "prestashop": ["prestashop"],
        }
        
        # Default to using the CMS name directly if not in our mapping
        search_terms = cms_keywords.get(cms_name.lower(), [cms_name.lower()])
        
        # List to store all vulnerabilities found
        all_vulnerabilities = []
        
        # Search for each term
        for term in search_terms:
            vulnerabilities = self._search_nvd_api(term, version)
            if vulnerabilities:
                all_vulnerabilities.extend(vulnerabilities)
                
        # Filter out duplicates by CVE ID
        deduplicated = {}
        for vuln in all_vulnerabilities:
            if vuln["cve_id"] not in deduplicated:
                deduplicated[vuln["cve_id"]] = vuln
                
        result = list(deduplicated.values())
        
        # Cache the result
        self.cache[cache_key] = result
        
        return result
        
    def check_webserver_vulnerabilities(self, server_name: str, version: str) -> List[Dict[str, Any]]:
        """Check vulnerabilities for a web server version.
        
        Args:
            server_name (str): Name of the web server (e.g., "Apache", "Nginx").
            version (str): Version number of the web server.
            
        Returns:
            List[Dict[str, Any]]: List of vulnerability information.
        """
        if not server_name or not version:
            return []
            
        # Create a cache key for this query
        cache_key = f"server-{server_name.lower()}-{version}"
        
        # Check cache first
        if cache_key in self.cache:
            self.logger.info(f"Using cached vulnerability data for {server_name} {version}")
            return self.cache[cache_key]
            
        # Normalize server name for searching
        server_keywords = {
            "apache": ["apache", "httpd"],
            "nginx": ["nginx"],
            "iis": ["iis", "microsoft iis", "internet information services"],
            "lightspeed": ["lightspeed"],
            "litespeed": ["litespeed"],
            "lighthttp": ["lighthttp"],
            "tomcat": ["tomcat", "apache tomcat"],
        }
        
        # Default to using the server name directly if not in our mapping
        search_terms = server_keywords.get(server_name.lower(), [server_name.lower()])
        
        # List to store all vulnerabilities found
        all_vulnerabilities = []
        
        # Search for each term
        for term in search_terms:
            vulnerabilities = self._search_nvd_api(term, version)
            if vulnerabilities:
                all_vulnerabilities.extend(vulnerabilities)
                
        # Filter out duplicates by CVE ID
        deduplicated = {}
        for vuln in all_vulnerabilities:
            if vuln["cve_id"] not in deduplicated:
                deduplicated[vuln["cve_id"]] = vuln
                
        result = list(deduplicated.values())
        
        # Cache the result
        self.cache[cache_key] = result
        
        return result
        
    def get_vulnerability_summary(self, vulnerabilities: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Generate a summary of vulnerabilities.
        
        Args:
            vulnerabilities (List[Dict[str, Any]]): List of vulnerability information.
            
        Returns:
            Dict[str, Any]: Summary information.
        """
        if not vulnerabilities:
            return self.create_empty_summary()
            
        # Count by severity
        severity_counts = {
            "critical": 0,
            "high": 0,
            "medium": 0,
            "low": 0,
            "unknown": 0,
        }
        
        exploitable_count = 0
        
        for vuln in vulnerabilities:
            severity = vuln.get("severity", "unknown")
            severity_counts[severity] = severity_counts.get(severity, 0) + 1
            
            if vuln.get("exploitable", False):
                exploitable_count += 1
                
        # Determine highest severity
        highest_severity = "none"
        for severity in ["critical", "high", "medium", "low"]:
            if severity_counts.get(severity, 0) > 0:
                highest_severity = severity
                break
                
        return {
            "total": len(vulnerabilities),
            "critical": severity_counts.get("critical", 0),
            "high": severity_counts.get("high", 0),
            "medium": severity_counts.get("medium", 0),
            "low": severity_counts.get("low", 0),
            "unknown": severity_counts.get("unknown", 0),
            "exploitable": exploitable_count,
            "has_vulnerabilities": len(vulnerabilities) > 0,
            "highest_severity": highest_severity,
        }
    
    def _search_nvd_api(self, keyword: str, version: str) -> List[Dict[str, Any]]:
        """Search the NVD API for vulnerabilities.
        
        Args:
            keyword (str): Keyword to search (e.g., "wordpress").
            version (str): Version number to check.
            
        Returns:
            List[Dict[str, Any]]: List of vulnerability information.
        """
        # Build the search query
        # Use CPE criteria for more precise matching when possible
        params = {
            "keywordSearch": f"{keyword} {version}",
            "resultsPerPage": 50,
        }
        
        headers = {}
        if self.api_key:
            headers["apiKey"] = self.api_key
        
        retries = 0
        while retries < self.max_retries:
            try:
                self.logger.debug(f"Querying NVD API for {keyword} {version}")
                response = requests.get(
                    self.api_base_url,
                    params=params,
                    headers=headers,
                    timeout=15
                )
                
                # Handle rate limiting
                if response.status_code == 429:
                    wait_time = int(response.headers.get("Retry-After", self.retry_delay))
                    self.logger.warning(f"Rate limited by NVD API. Waiting {wait_time} seconds.")
                    time.sleep(wait_time)
                    retries += 1
                    continue
                    
                # Handle other errors
                if response.status_code != 200:
                    self.logger.error(f"NVD API error: {response.status_code} - {response.text}")
                    return []
                    
                data = response.json()
                
                # Process the results
                vulnerabilities = []
                
                for vuln in data.get("vulnerabilities", []):
                    cve = vuln.get("cve", {})
                    
                    # Check if this vulnerability affects the specified version
                    if not self._check_version_affected(cve, keyword, version):
                        continue
                        
                    # Extract vulnerability information
                    vulnerability = {
                        "cve_id": cve.get("id", "Unknown"),
                        "description": self._get_cve_description(cve),
                        "severity": self._get_cve_severity(cve),
                        "cvss_score": self._get_cvss_score(cve),
                        "published_date": cve.get("published", ""),
                        "last_modified": cve.get("lastModified", ""),
                        "references": self._get_cve_references(cve),
                        "exploitable": self._is_exploitable(cve),
                    }
                    
                    vulnerabilities.append(vulnerability)
                
                return vulnerabilities
                
            except Exception as e:
                self.logger.error(f"Error searching NVD API: {str(e)}")
                retries += 1
                time.sleep(self.retry_delay)
        
        self.logger.error(f"Failed to query NVD API after {self.max_retries} attempts")
        return []
        
    def _check_version_affected(self, cve: Dict[str, Any], keyword: str, version: str) -> bool:
        """Check if a CVE affects the specified version.
        
        Args:
            cve (Dict[str, Any]): CVE data from NVD API.
            keyword (str): Keyword used in search.
            version (str): Version to check.
            
        Returns:
            bool: True if the version is affected, False otherwise.
        """
        # Extract version components for comparison
        version_parts = version.split('.')
        version_nums = [int(part) for part in version_parts if part.isdigit()]
        
        # Pad version with zeros for consistent comparison
        while len(version_nums) < 3:
            version_nums.append(0)
            
        # Check configurations
        configurations = cve.get("configurations", [])
        
        for config in configurations:
            nodes = config.get("nodes", [])
            
            for node in nodes:
                matches = node.get("cpeMatch", [])
                
                for match in matches:
                    criteria = match.get("criteria", "").lower()
                    
                    # If criteria contains our keyword
                    if keyword.lower() in criteria:
                        # Check version constraints
                        version_start_including = match.get("versionStartIncluding", "")
                        version_start_excluding = match.get("versionStartExcluding", "")
                        version_end_including = match.get("versionEndIncluding", "")
                        version_end_excluding = match.get("versionEndExcluding", "")
                        
                        # If no version constraints, consider it matching
                        if not any([version_start_including, version_start_excluding, 
                                   version_end_including, version_end_excluding]):
                            return True
                            
                        # Check version range
                        is_affected = True
                        
                        if version_start_including and not self._version_greater_or_equal(version, version_start_including):
                            is_affected = False
                            
                        if version_start_excluding and not self._version_greater_than(version, version_start_excluding):
                            is_affected = False
                            
                        if version_end_including and not self._version_less_or_equal(version, version_end_including):
                            is_affected = False
                            
                        if version_end_excluding and not self._version_less_than(version, version_end_excluding):
                            is_affected = False
                            
                        if is_affected:
                            return True
        
        # Also check if the vulnerability description mentions the exact version
        description = self._get_cve_description(cve).lower()
        
        # Check for exact version match or version range descriptions
        version_indicators = [
            f"{keyword} {version}",
            f"{keyword} version {version}",
            f"{keyword} versions {version}",
            f"affecting {keyword} {version}",
            f"affects {keyword} {version}",
            f"vulnerable {keyword} {version}",
        ]
        
        for indicator in version_indicators:
            if indicator.lower() in description:
                return True
        
        # More sophisticated version range check in text
        # This is a simple approach and might need refinement
        if (f"versions prior to" in description or 
            f"versions before" in description or 
            f"versions up to" in description or 
            f"versions earlier than" in description):
            # Extracting version number after these phrases is complex and prone to errors
            # This is a simplified check
            if f"{version}" in description:
                return True
        
        return False
        
    def _version_greater_than(self, version1: str, version2: str) -> bool:
        """Check if version1 is greater than version2.
        
        Args:
            version1 (str): First version to compare.
            version2 (str): Second version to compare.
            
        Returns:
            bool: True if version1 > version2, False otherwise.
        """
        v1_parts = [int(p) for p in version1.split('.') if p.isdigit()]
        v2_parts = [int(p) for p in version2.split('.') if p.isdigit()]
        
        # Pad with zeros for consistent comparison
        while len(v1_parts) < len(v2_parts):
            v1_parts.append(0)
        while len(v2_parts) < len(v1_parts):
            v2_parts.append(0)
            
        for i in range(len(v1_parts)):
            if v1_parts[i] > v2_parts[i]:
                return True
            if v1_parts[i] < v2_parts[i]:
                return False
                
        return False  # Versions are equal
        
    def _version_greater_or_equal(self, version1: str, version2: str) -> bool:
        """Check if version1 is greater than or equal to version2.
        
        Args:
            version1 (str): First version to compare.
            version2 (str): Second version to compare.
            
        Returns:
            bool: True if version1 >= version2, False otherwise.
        """
        if version1 == version2:
            return True
        return self._version_greater_than(version1, version2)
        
    def _version_less_than(self, version1: str, version2: str) -> bool:
        """Check if version1 is less than version2.
        
        Args:
            version1 (str): First version to compare.
            version2 (str): Second version to compare.
            
        Returns:
            bool: True if version1 < version2, False otherwise.
        """
        return not self._version_greater_or_equal(version1, version2)
        
    def _version_less_or_equal(self, version1: str, version2: str) -> bool:
        """Check if version1 is less than or equal to version2.
        
        Args:
            version1 (str): First version to compare.
            version2 (str): Second version to compare.
            
        Returns:
            bool: True if version1 <= version2, False otherwise.
        """
        return not self._version_greater_than(version1, version2)
        
    def _get_cve_description(self, cve: Dict[str, Any]) -> str:
        """Extract the English description from a CVE.
        
        Args:
            cve (Dict[str, Any]): CVE data.
            
        Returns:
            str: Description text or empty string if not found.
        """
        descriptions = cve.get("descriptions", [])
        
        for desc in descriptions:
            if desc.get("lang") == "en":
                return desc.get("value", "")
                
        return ""
        
    def _get_cve_severity(self, cve: Dict[str, Any]) -> str:
        """Extract the severity rating from a CVE.
        
        Args:
            cve (Dict[str, Any]): CVE data.
            
        Returns:
            str: Severity rating (critical, high, medium, low, or unknown).
        """
        metrics = cve.get("metrics", {})
        
        # Check CVSS v3.1 first
        cvss_v31 = metrics.get("cvssMetricV31", [])
        if cvss_v31 and len(cvss_v31) > 0:
            base_score = cvss_v31[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return self._score_to_severity(base_score)
                
        # Fall back to CVSS v3.0
        cvss_v30 = metrics.get("cvssMetricV30", [])
        if cvss_v30 and len(cvss_v30) > 0:
            base_score = cvss_v30[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return self._score_to_severity(base_score)
                
        # Fall back to CVSS v2.0
        cvss_v2 = metrics.get("cvssMetricV2", [])
        if cvss_v2 and len(cvss_v2) > 0:
            base_score = cvss_v2[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return self._score_to_severity(base_score)
                
        return "unknown"
        
    def _get_cvss_score(self, cve: Dict[str, Any]) -> Optional[float]:
        """Extract the CVSS base score from a CVE.
        
        Args:
            cve (Dict[str, Any]): CVE data.
            
        Returns:
            Optional[float]: CVSS base score or None if not found.
        """
        metrics = cve.get("metrics", {})
        
        # Check CVSS v3.1 first
        cvss_v31 = metrics.get("cvssMetricV31", [])
        if cvss_v31 and len(cvss_v31) > 0:
            base_score = cvss_v31[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return base_score
                
        # Fall back to CVSS v3.0
        cvss_v30 = metrics.get("cvssMetricV30", [])
        if cvss_v30 and len(cvss_v30) > 0:
            base_score = cvss_v30[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return base_score
                
        # Fall back to CVSS v2.0
        cvss_v2 = metrics.get("cvssMetricV2", [])
        if cvss_v2 and len(cvss_v2) > 0:
            base_score = cvss_v2[0].get("cvssData", {}).get("baseScore")
            if base_score is not None:
                return base_score
                
        return None
        
    def _score_to_severity(self, score: float) -> str:
        """Convert a CVSS score to a severity rating.
        
        Args:
            score (float): CVSS base score.
            
        Returns:
            str: Severity rating (critical, high, medium, or low).
        """
        if score >= 9.0:
            return "critical"
        elif score >= 7.0:
            return "high"
        elif score >= 4.0:
            return "medium"
        else:
            return "low"
            
    def _get_cve_references(self, cve: Dict[str, Any]) -> List[str]:
        """Extract reference URLs from a CVE.
        
        Args:
            cve (Dict[str, Any]): CVE data.
            
        Returns:
            List[str]: List of reference URLs.
        """
        references = []
        
        for ref in cve.get("references", []):
            url = ref.get("url")
            if url:
                references.append(url)
                
        return references
        
    def _is_exploitable(self, cve: Dict[str, Any]) -> bool:
        """Check if a CVE has known exploits.
        
        Args:
            cve (Dict[str, Any]): CVE data.
            
        Returns:
            bool: True if the CVE has known exploits, False otherwise.
        """
        # Check for exploit references
        for ref in cve.get("references", []):
            tags = ref.get("tags", [])
            
            if "Exploit" in tags or "exploit" in tags or "Malware" in tags:
                return True
                
        # Check vulnerability descriptions for exploit mentions
        description = self._get_cve_description(cve).lower()
        
        exploit_indicators = [
            "exploit",
            "exploited",
            "exploitable",
            "in the wild",
            "working exploit",
            "proof of concept",
            "poc available",
        ]
        
        for indicator in exploit_indicators:
            if indicator in description:
                return True
                
        return False