"""Vulnerability scanner plugin.

Runs one or more vulnerability providers against an AiSbomDocument and
returns a unified list of findings consumed by other plugins (e.g.
``sarif_export``, ``cyclonedx_export``).

Providers
---------
vela-rules
    Built-in structural security checks on SBOM nodes — no network required.
    Checks for missing authentication, overly-permissive tool access,
    unencrypted endpoints, and other AI-stack risks.  Always fast.
osv
    Queries https://api.osv.dev for known CVEs in the SBOM's ``deps`` list.
    Requires network access.
grype
    Runs the Grype binary (must be on ``PATH``) against the deps list
    exported as a minimal CycloneDX BOM.  Requires Grype ≥ 0.60.
all
    Runs all three providers in sequence and merges findings.

Config keys
-----------
provider         vela-rules | osv | grype | all   (default: vela-rules)
timeout          OSV network timeout in seconds    (default: 15.0)
grype_timeout    Grype subprocess timeout          (default: 60.0)
"""
from __future__ import annotations

import json
import logging
import re
import shutil
import subprocess
import tempfile
from typing import Any

from ..models import ToolResult
from ..plugin_base import ToolPlugin

_log = logging.getLogger("toolbox.plugins.vulnerability")

# ── Finding schema ────────────────────────────────────────────────────────────
#
#  rule_id       str   — unique identifier (e.g. NGA-001, CVE-2024-…)
#  title         str   — short human-readable title
#  description   str   — full description
#  severity      str   — CRITICAL | HIGH | MEDIUM | LOW | INFO
#  affected      list  — affected node IDs, package PURLs, or component names
#  remediation   str   — suggested fix
#  advisory_url  str   — link to more information (may be empty string)
#
# ─────────────────────────────────────────────────────────────────────────────

_SEVERITY_RANK = {"CRITICAL": 4, "HIGH": 3, "MEDIUM": 2, "LOW": 1, "INFO": 0}


# ── Vela/NGA structural rules ─────────────────────────────────────────────────

def _vela_rules_scan(sbom: dict[str, Any]) -> list[dict[str, Any]]:
    """Run built-in structural checks over SBOM nodes.

    Each check maps to a rule ID in the NGA-V (Vulnerability) sub-family.
    """
    findings: list[dict[str, Any]] = []
    nodes: list[dict[str, Any]] = sbom.get("nodes") or []
    edges: list[dict[str, Any]] = sbom.get("edges") or []

    # Build lookup: node_id → node
    node_map: dict[str, dict[str, Any]] = {
        n.get("node_id", ""): n for n in nodes if n.get("node_id")
    }

    # ── NGA-V001: unauthenticated agent/tool endpoint ─────────────────────
    for node in nodes:
        ctype = (node.get("component_type") or "").upper()
        if ctype not in {"AGENT", "TOOL", "API_ENDPOINT"}:
            continue
        meta = node.get("metadata") or {}
        risk_tags: list[str] = meta.get("risk_tags") or []
        if any("no-auth" in t.lower() for t in risk_tags):
            findings.append(
                {
                    "rule_id":      "NGA-V001",
                    "title":        "Unauthenticated AI component",
                    "description":  (
                        f"Component '{node.get('name', node.get('node_id', ''))}' "
                        f"is tagged 'no-auth-required', making it accessible without "
                        f"authentication. This can enable prompt injection, data "
                        f"exfiltration, or unauthorised use."
                    ),
                    "severity":     "HIGH",
                    "affected":     [node.get("node_id", node.get("name", ""))],
                    "remediation":  (
                        "Add authentication middleware (API key, OAuth2, or mTLS) "
                        "in front of this component."
                    ),
                    "advisory_url": "https://owasp.org/www-project-top-10-for-large-language-model-applications/",
                }
            )

    # ── NGA-V002: tool with shell/exec permissions ────────────────────────
    _EXEC_PERMS = {"shell", "exec", "execute", "subprocess", "os.system", "eval"}
    for node in nodes:
        ctype = (node.get("component_type") or "").upper()
        if ctype != "TOOL":
            continue
        meta = node.get("metadata") or {}
        perms: list[str] = [p.lower() for p in (meta.get("permissions") or [])]
        exec_perms = [p for p in perms if any(ep in p for ep in _EXEC_PERMS)]
        if exec_perms:
            findings.append(
                {
                    "rule_id":      "NGA-V002",
                    "title":        "Tool has shell/exec permissions",
                    "description":  (
                        f"Tool '{node.get('name', node.get('node_id', ''))}' "
                        f"has execution permissions ({', '.join(exec_perms)}). "
                        f"LLM agents with access to this tool can run arbitrary "
                        f"commands on the host."
                    ),
                    "severity":     "CRITICAL",
                    "affected":     [node.get("node_id", node.get("name", ""))],
                    "remediation":  (
                        "Restrict tool permissions to least-privilege. Consider "
                        "sandboxing command execution with containers or gVisor."
                    ),
                    "advisory_url": "https://atlas.mitre.org/techniques/AML.T0052/",
                }
            )

    # ── NGA-V003: plaintext (http://) endpoint ────────────────────────────
    for node in nodes:
        url = node.get("url") or ""
        if re.match(r"http://", url, re.IGNORECASE):
            findings.append(
                {
                    "rule_id":      "NGA-V003",
                    "title":        "Unencrypted HTTP endpoint",
                    "description":  (
                        f"Node '{node.get('name', node.get('node_id', ''))}' "
                        f"communicates over plain HTTP ({url}). "
                        f"Prompt payloads and LLM responses are transmitted in "
                        f"cleartext, enabling man-in-the-middle attacks."
                    ),
                    "severity":     "HIGH",
                    "affected":     [node.get("node_id", node.get("name", ""))],
                    "remediation":  "Migrate to HTTPS/TLS. Use HTTP Strict-Transport-Security.",
                    "advisory_url": "",
                }
            )

    # ── NGA-V004: SQL-injectable tool ─────────────────────────────────────
    for node in nodes:
        meta = node.get("metadata") or {}
        risk_tags: list[str] = meta.get("risk_tags") or []
        if any("sql-inject" in t.lower() for t in risk_tags):
            findings.append(
                {
                    "rule_id":      "NGA-V004",
                    "title":        "SQL-injectable tool parameter",
                    "description":  (
                        f"Tool '{node.get('name', node.get('node_id', ''))}' "
                        f"is tagged 'SQL-injectable'. An LLM that processes "
                        f"untrusted user input could be prompted to inject SQL "
                        f"through this tool."
                    ),
                    "severity":     "HIGH",
                    "affected":     [node.get("node_id", node.get("name", ""))],
                    "remediation":  (
                        "Use parameterised queries or an ORM. Validate and sanitise "
                        "all inputs before passing to SQL tools."
                    ),
                    "advisory_url": "https://owasp.org/www-community/attacks/SQL_Injection",
                }
            )

    # ── NGA-V005: agent with WRITE access to sensitive datastore ──────────
    _SENSITIVE_TYPES = {"DATABASE", "DATASTORE", "VECTOR_STORE", "KNOWLEDGE_BASE"}
    for edge in edges:
        if (edge.get("relationship") or "").upper() != "ACCESSES":
            continue
        if (edge.get("access_type") or "").lower() not in ("write", "readwrite"):
            continue
        target_id = edge.get("target_id") or edge.get("target") or ""
        target = node_map.get(target_id)
        if target is None:
            continue
        if (target.get("component_type") or "").upper() not in _SENSITIVE_TYPES:
            continue
        source_id = edge.get("source_id") or edge.get("source") or ""
        source = node_map.get(source_id)
        source_name = (source.get("name") if source else None) or source_id
        target_name = target.get("name") or target_id
        findings.append(
            {
                "rule_id":      "NGA-V005",
                "title":        "Agent has write access to sensitive datastore",
                "description":  (
                    f"'{source_name}' has write/readwrite access to "
                    f"'{target_name}' (type: {target.get('component_type', '')}). "
                    f"A compromised or prompt-injected agent could corrupt, "
                    f"exfiltrate, or poison the data store."
                ),
                "severity":     "MEDIUM",
                "affected":     [source_id, target_id],
                "remediation":  (
                    "Apply least-privilege: grant read-only access unless writes "
                    "are explicitly required. Log all write operations."
                ),
                "advisory_url": "https://atlas.mitre.org/techniques/AML.T0048/",
            }
        )

    # ── NGA-V006: MCP server exposed without auth ─────────────────────────
    for node in nodes:
        ctype = (node.get("component_type") or "").upper()
        name_lower = (node.get("name") or "").lower()
        is_mcp = ctype == "MCP_SERVER" or "mcp" in name_lower
        if not is_mcp:
            continue
        meta = node.get("metadata") or {}
        risk_tags = meta.get("risk_tags") or []
        if any("no-auth" in t.lower() for t in risk_tags):
            findings.append(
                {
                    "rule_id":      "NGA-V006",
                    "title":        "MCP server exposed without authentication",
                    "description":  (
                        f"MCP server '{node.get('name', node.get('node_id', ''))}' "
                        f"is reachable without authentication. Any client (including "
                        f"malicious ones) can call its tools, invoke resources, and "
                        f"potentially exfiltrate data via the MCP protocol."
                    ),
                    "severity":     "HIGH",
                    "affected":     [node.get("node_id", node.get("name", ""))],
                    "remediation":  (
                        "Enable OAuth 2.0 / bearer token auth on the MCP server. "
                        "Restrict the listening interface to localhost or a VPN."
                    ),
                    "advisory_url": "https://owasp.org/www-project-top-10-for-large-language-model-applications/",
                }
            )

    return findings


# ── OSV provider ──────────────────────────────────────────────────────────────

def _osv_scan(
    sbom: dict[str, Any], timeout: float
) -> list[dict[str, Any]]:
    """Translate OSV results into the shared finding schema."""
    try:
        from ..osv_client import query_osv
    except ImportError:
        _log.warning("osv_client not available; skipping OSV scan")
        return []

    deps = sbom.get("deps") or []
    if not deps:
        _log.info("no deps in SBOM; skipping OSV scan")
        return []

    osv_results = query_osv(deps, timeout=timeout)
    findings: list[dict[str, Any]] = []
    for osv in osv_results:
        advisory_id = osv["advisory_id"]
        cve_ids = osv.get("cve_ids") or []
        rule_id = cve_ids[0] if cve_ids else advisory_id
        dep_name = osv["dep_name"]
        affected_str = osv.get("affected_versions") or "see advisory"
        findings.append(
            {
                "rule_id":      rule_id,
                "title":        f"{advisory_id} in {dep_name}",
                "description":  osv.get("summary") or f"Known vulnerability in {dep_name}.",
                "severity":     osv.get("severity", "UNKNOWN"),
                "affected":     [osv.get("purl") or dep_name],
                "remediation":  (
                    f"Upgrade {dep_name} to a version outside "
                    f"the affected range ({affected_str})."
                ),
                "advisory_url": osv.get("url") or "",
            }
        )
    return findings


# ── Grype provider ────────────────────────────────────────────────────────────

def _grype_scan(
    sbom: dict[str, Any], grype_timeout: float
) -> list[dict[str, Any]]:
    """Run Grype against a minimal CycloneDX BOM and parse its JSON output."""
    if not shutil.which("grype"):
        _log.info("grype not on PATH; skipping grype scan")
        return []

    deps = sbom.get("deps") or []
    if not deps:
        return []

    # Build a minimal CycloneDX 1.4 BOM for Grype
    components = []
    for dep in deps:
        purl = dep.get("purl") or ""
        name = dep.get("name") or purl
        if not purl and not name:
            continue
        comp: dict[str, Any] = {"type": "library", "name": name}
        if purl:
            comp["purl"] = purl
        # Extract version from PURL or version_spec
        version_spec = dep.get("version_spec") or ""
        pin = re.match(r"==\s*([\w.\-+]+)", version_spec)
        if pin:
            comp["version"] = pin.group(1)
        else:
            m = re.search(r"@([\w.\-+]+)$", purl)
            if m:
                comp["version"] = m.group(1)
        components.append(comp)

    if not components:
        return []

    bom = {
        "bomFormat": "CycloneDX",
        "specVersion": "1.4",
        "version": 1,
        "components": components,
    }

    findings: list[dict[str, Any]] = []
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".cdx.json", delete=False
        ) as fh:
            json.dump(bom, fh)
            bom_path = fh.name

        result = subprocess.run(  # noqa: S603
            ["grype", f"sbom:{bom_path}", "-o", "json", "--quiet"],
            capture_output=True,
            text=True,
            timeout=grype_timeout,
            check=False,
        )
        if result.returncode not in (0, 1):
            _log.warning("grype exited %d: %s", result.returncode, result.stderr[:300])
            return []

        data = json.loads(result.stdout)
        for match in data.get("matches") or []:
            vuln = match.get("vulnerability") or {}
            artifact = match.get("artifact") or {}
            cve_id = vuln.get("id", "")
            severity = (vuln.get("severity") or "UNKNOWN").upper()
            if severity not in _SEVERITY_RANK:
                severity = "UNKNOWN"
            fix_info = vuln.get("fix") or {}
            fix_versions = fix_info.get("versions") or []
            fix_str = (
                f"Upgrade to {', '.join(fix_versions)}"
                if fix_versions
                else "See advisory for remediation."
            )
            pkg_name = artifact.get("name") or ""
            purl = artifact.get("purl") or pkg_name
            findings.append(
                {
                    "rule_id":      cve_id or "GRYPE-UNKNOWN",
                    "title":        f"{cve_id} in {pkg_name}" if cve_id else f"Vulnerability in {pkg_name}",
                    "description":  vuln.get("description") or vuln.get("id") or "",
                    "severity":     severity,
                    "affected":     [purl],
                    "remediation":  fix_str,
                    "advisory_url": (vuln.get("dataSource") or ""),
                }
            )
    except subprocess.TimeoutExpired:
        _log.warning("grype timed out after %.0fs", grype_timeout)
    except (json.JSONDecodeError, OSError, KeyError) as exc:
        _log.warning("grype output parse error: %s", exc)

    _log.info("grype scan: %d finding(s)", len(findings))
    return findings


# ── Plugin ────────────────────────────────────────────────────────────────────

class VulnerabilityScannerPlugin(ToolPlugin):
    """Run one or more vulnerability providers and return a unified finding list.

    Config keys
    -----------
    provider       vela-rules | osv | grype | all   (default: vela-rules)
    timeout        OSV network timeout in seconds    (default: 15.0)
    grype_timeout  Grype subprocess timeout          (default: 60.0)
    """

    name = "vulnerability"

    def run(self, sbom: dict[str, Any], config: dict[str, Any]) -> ToolResult:
        provider = config.get("provider", "vela-rules")
        timeout = float(config.get("timeout", 15.0))
        grype_timeout = float(config.get("grype_timeout", 60.0))

        all_findings: list[dict[str, Any]] = []

        if provider in ("vela-rules", "all"):
            vr = _vela_rules_scan(sbom)
            _log.info("vela-rules: %d finding(s)", len(vr))
            all_findings.extend(vr)

        if provider in ("osv", "all"):
            osv = _osv_scan(sbom, timeout)
            all_findings.extend(osv)

        if provider in ("grype", "all"):
            grype = _grype_scan(sbom, grype_timeout)
            all_findings.extend(grype)

        if provider not in ("vela-rules", "osv", "grype", "all"):
            _log.warning("unknown provider '%s'; falling back to vela-rules", provider)
            all_findings.extend(_vela_rules_scan(sbom))

        # Sort by descending severity
        all_findings.sort(
            key=lambda f: _SEVERITY_RANK.get(f.get("severity", "INFO"), 0),
            reverse=True,
        )

        total = len(all_findings)
        status = "warning" if total else "ok"
        return ToolResult(
            status=status,
            tool=self.name,
            message=f"{total} vulnerability finding(s) (provider={provider})",
            details={"findings": all_findings},
        )
