"""
MundiX Threat Intelligence Module
==================================

Proactively queries external threat intelligence sources to enrich
engagement findings and help the agent prioritize exploitation.

Sources:
  - NVD (NIST National Vulnerability Database) API 2.0
  - ExploitDB via searchsploit (local)
  - CISA KEV (Known Exploited Vulnerabilities) catalog

Usage:
    from threat_intel import ThreatIntel
    ti = ThreatIntel()
    
    # Enrich a service with CVE data
    cves = ti.search_nvd("Apache", "2.4.41")
    
    # Check if a CVE is actively exploited
    kev = ti.check_cisa_kev("CVE-2024-12345")
    
    # Search for exploits
    exploits = ti.search_exploitdb("Apache 2.4.41")
    
    # Full enrichment from discovered ports
    enriched = ti.enrich_ports([
        {"port": 80, "service": "http", "version": "Apache/2.4.41"},
        {"port": 22, "service": "ssh", "version": "OpenSSH 8.2"},
    ])
"""

import json
import os
import subprocess
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple

try:
    import requests
except ImportError:
    requests = None


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

NVD_API_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0"
CISA_KEV_URL = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"

# Cache directory
CACHE_DIR = os.path.expanduser("~/.mundix/threat_intel_cache")
os.makedirs(CACHE_DIR, exist_ok=True)

# Cache TTL
NVD_CACHE_TTL = 3600 * 6       # 6 hours
KEV_CACHE_TTL = 3600 * 24      # 24 hours
EXPLOITDB_CACHE_TTL = 3600 * 2  # 2 hours

# Rate limiting for NVD (6 req/min without API key, 50 with)
NVD_RATE_LIMIT = 10  # seconds between requests


class ThreatIntel:
    """Threat intelligence aggregator for pentest enrichment."""

    def __init__(self, nvd_api_key: str = None):
        self.nvd_api_key = nvd_api_key or os.environ.get("NVD_API_KEY")
        self._last_nvd_call = 0
        self._kev_catalog = None
        self._kev_loaded_at = 0

    # --- NVD (National Vulnerability Database) ---

    def search_nvd(self, product: str, version: str = None,
                   max_results: int = 10) -> List[Dict]:
        """
        Search NVD for CVEs related to a product/version.
        Returns list of CVE dicts with id, description, cvss, severity, references.
        """
        if not requests:
            return self._search_nvd_curl(product, version, max_results)

        # Check cache
        cache_key = f"nvd_{product}_{version or 'any'}".replace(" ", "_").lower()
        cached = self._read_cache(cache_key, NVD_CACHE_TTL)
        if cached is not None:
            return cached[:max_results]

        # Rate limit
        self._nvd_rate_limit()

        keyword = f"{product} {version}" if version else product
        params = {
            "keywordSearch": keyword,
            "resultsPerPage": max_results,
        }
        headers = {}
        if self.nvd_api_key:
            headers["apiKey"] = self.nvd_api_key

        try:
            resp = requests.get(NVD_API_URL, params=params, headers=headers, timeout=30)
            resp.raise_for_status()
            data = resp.json()

            cves = []
            for item in data.get("vulnerabilities", []):
                cve = item.get("cve", {})
                cve_id = cve.get("id", "")

                # Get CVSS score
                cvss_score = 0.0
                cvss_severity = "unknown"
                metrics = cve.get("metrics", {})
                for cvss_key in ["cvssMetricV31", "cvssMetricV30", "cvssMetricV2"]:
                    if cvss_key in metrics:
                        cvss_data = metrics[cvss_key][0].get("cvssData", {})
                        cvss_score = cvss_data.get("baseScore", 0.0)
                        cvss_severity = cvss_data.get("baseSeverity", "UNKNOWN").lower()
                        break

                # Get description
                descriptions = cve.get("descriptions", [])
                desc = next((d["value"] for d in descriptions if d["lang"] == "en"), "")

                # Get references
                refs = [r["url"] for r in cve.get("references", [])[:3]]

                cves.append({
                    "id": cve_id,
                    "description": desc[:300],
                    "cvss_score": cvss_score,
                    "severity": cvss_severity,
                    "references": refs,
                    "published": cve.get("published", ""),
                    "source": "nvd",
                })

            # Sort by CVSS score descending
            cves.sort(key=lambda x: x["cvss_score"], reverse=True)
            self._write_cache(cache_key, cves)
            return cves[:max_results]

        except Exception as e:
            return [{"error": str(e), "source": "nvd"}]

    def _search_nvd_curl(self, product: str, version: str = None,
                         max_results: int = 10) -> List[Dict]:
        """Fallback: use curl if requests is not available."""
        keyword = f"{product} {version}" if version else product
        try:
            cmd = (
                f'curl -s "{NVD_API_URL}?keywordSearch={keyword}'
                f'&resultsPerPage={max_results}"'
            )
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            if result.returncode == 0:
                data = json.loads(result.stdout)
                cves = []
                for item in data.get("vulnerabilities", []):
                    cve = item.get("cve", {})
                    cves.append({
                        "id": cve.get("id", ""),
                        "description": str(cve.get("descriptions", [{}])[0].get("value", ""))[:300],
                        "source": "nvd",
                    })
                return cves
        except Exception:
            pass
        return []

    # --- CISA KEV (Known Exploited Vulnerabilities) ---

    def load_cisa_kev(self) -> Dict:
        """Load the CISA KEV catalog (cached locally)."""
        now = time.time()
        if self._kev_catalog and (now - self._kev_loaded_at) < KEV_CACHE_TTL:
            return self._kev_catalog

        # Try cache
        cached = self._read_cache("cisa_kev", KEV_CACHE_TTL)
        if cached is not None:
            self._kev_catalog = cached
            self._kev_loaded_at = now
            return cached

        # Download
        try:
            if requests:
                resp = requests.get(CISA_KEV_URL, timeout=30)
                resp.raise_for_status()
                data = resp.json()
            else:
                result = subprocess.run(
                    f'curl -s "{CISA_KEV_URL}"',
                    shell=True, capture_output=True, text=True, timeout=30
                )
                data = json.loads(result.stdout)

            # Index by CVE ID for fast lookup
            catalog = {}
            for vuln in data.get("vulnerabilities", []):
                cve_id = vuln.get("cveID", "")
                catalog[cve_id] = {
                    "vendor": vuln.get("vendorProject", ""),
                    "product": vuln.get("product", ""),
                    "name": vuln.get("vulnerabilityName", ""),
                    "date_added": vuln.get("dateAdded", ""),
                    "due_date": vuln.get("dueDate", ""),
                    "known_ransomware": vuln.get("knownRansomwareCampaignUse", "Unknown"),
                    "notes": vuln.get("notes", ""),
                }

            self._kev_catalog = catalog
            self._kev_loaded_at = now
            self._write_cache("cisa_kev", catalog)
            return catalog

        except Exception:
            return {}

    def check_cisa_kev(self, cve_id: str) -> Optional[Dict]:
        """Check if a CVE is in the CISA Known Exploited Vulnerabilities catalog."""
        catalog = self.load_cisa_kev()
        return catalog.get(cve_id)

    def get_critical_kevs(self, product: str = None) -> List[Dict]:
        """Get all KEVs, optionally filtered by product name."""
        catalog = self.load_cisa_kev()
        results = []
        for cve_id, info in catalog.items():
            if product and product.lower() not in info.get("product", "").lower():
                continue
            results.append({"cve_id": cve_id, **info})
        return results

    # --- ExploitDB (via searchsploit) ---

    def search_exploitdb(self, query: str, max_results: int = 10) -> List[Dict]:
        """Search ExploitDB for known exploits using searchsploit."""
        cache_key = f"exploitdb_{query}".replace(" ", "_").lower()
        cached = self._read_cache(cache_key, EXPLOITDB_CACHE_TTL)
        if cached is not None:
            return cached[:max_results]

        try:
            result = subprocess.run(
                ["searchsploit", query, "--json"],
                capture_output=True, text=True, timeout=30
            )
            if result.returncode == 0 and result.stdout.strip():
                data = json.loads(result.stdout)
                exploits = []
                for exp in data.get("RESULTS_EXPLOIT", [])[:max_results]:
                    exploits.append({
                        "title": exp.get("Title", ""),
                        "edb_id": exp.get("EDB-ID", ""),
                        "path": exp.get("Path", ""),
                        "date": exp.get("Date Published", ""),
                        "platform": exp.get("Platform", ""),
                        "type": exp.get("Type", ""),
                        "source": "exploitdb",
                    })
                self._write_cache(cache_key, exploits)
                return exploits
        except FileNotFoundError:
            pass  # searchsploit not installed
        except Exception:
            pass
        return []

    # --- Enrichment ---

    def enrich_ports(self, ports: List[Dict]) -> List[Dict]:
        """
        Enrich discovered ports with threat intel.
        For each port with a known service/version, search NVD and ExploitDB.
        Returns enriched port dicts with 'cves' and 'exploits' fields.
        """
        enriched = []
        for port in ports:
            service = port.get("service", "")
            version = port.get("version", "")
            port_info = dict(port)
            port_info["cves"] = []
            port_info["exploits"] = []
            port_info["kev_match"] = False

            if service and version:
                # Search NVD
                cves = self.search_nvd(service, version, max_results=5)
                port_info["cves"] = [c for c in cves if "error" not in c]

                # Check if any CVE is in CISA KEV
                for cve in port_info["cves"]:
                    kev = self.check_cisa_kev(cve["id"])
                    if kev:
                        cve["actively_exploited"] = True
                        cve["kev_info"] = kev
                        port_info["kev_match"] = True

                # Search ExploitDB
                search_term = f"{service} {version}"
                port_info["exploits"] = self.search_exploitdb(search_term, max_results=3)

            enriched.append(port_info)
        return enriched

    def get_enrichment_summary(self, ports: List[Dict]) -> str:
        """Get a human-readable threat intel summary for the AI context."""
        enriched = self.enrich_ports(ports)
        lines = ["## Threat Intelligence Summary\n"]

        total_cves = 0
        total_exploits = 0
        critical_cves = []
        kev_matches = []

        for port in enriched:
            cves = port.get("cves", [])
            exploits = port.get("exploits", [])
            total_cves += len(cves)
            total_exploits += len(exploits)

            for cve in cves:
                if cve.get("cvss_score", 0) >= 9.0:
                    critical_cves.append(cve)
                if cve.get("actively_exploited"):
                    kev_matches.append(cve)

            if cves or exploits:
                svc = f"{port.get('service', '?')} {port.get('version', '')}"
                lines.append(f"**Port {port.get('port', '?')}** ({svc}):")
                for cve in cves[:3]:
                    marker = " 🔥 ACTIVELY EXPLOITED" if cve.get("actively_exploited") else ""
                    lines.append(
                        f"  - {cve['id']} (CVSS {cve.get('cvss_score', '?')}/"
                        f"{cve.get('severity', '?')}){marker}"
                    )
                for exp in exploits[:2]:
                    lines.append(f"  - ExploitDB: {exp.get('title', '?')} (EDB-{exp.get('edb_id', '?')})")
                lines.append("")

        lines.insert(1, f"Total: {total_cves} CVEs, {total_exploits} exploits, "
                        f"{len(kev_matches)} actively exploited (CISA KEV)\n")

        if critical_cves:
            lines.append("### ⚠️ CRITICAL (CVSS ≥ 9.0) — Prioritize exploitation:")
            for cve in critical_cves[:5]:
                lines.append(f"  - {cve['id']}: {cve.get('description', '')[:100]}")

        if kev_matches:
            lines.append("\n### 🔥 CISA KEV — Known to be actively exploited in the wild:")
            for cve in kev_matches[:5]:
                kev = cve.get("kev_info", {})
                lines.append(f"  - {cve['id']}: {kev.get('name', '')} "
                             f"(ransomware: {kev.get('known_ransomware', 'Unknown')})")

        return '\n'.join(lines)

    # --- Cache Helpers ---

    def _read_cache(self, key: str, ttl: int) -> Optional[any]:
        """Read from file cache if not expired."""
        path = os.path.join(CACHE_DIR, f"{key}.json")
        try:
            if os.path.exists(path):
                age = time.time() - os.path.getmtime(path)
                if age < ttl:
                    with open(path, 'r') as f:
                        return json.load(f)
        except Exception:
            pass
        return None

    def _write_cache(self, key: str, data: any):
        """Write to file cache."""
        path = os.path.join(CACHE_DIR, f"{key}.json")
        try:
            with open(path, 'w') as f:
                json.dump(data, f, ensure_ascii=False)
        except Exception:
            pass

    def _nvd_rate_limit(self):
        """Enforce NVD rate limits."""
        elapsed = time.time() - self._last_nvd_call
        if elapsed < NVD_RATE_LIMIT:
            time.sleep(NVD_RATE_LIMIT - elapsed)
        self._last_nvd_call = time.time()
