vulncheck/app/services/override_jobs.py

"""
In-memory job tracker for long-running CVSS-override runs.

The CISA Vulnrichment override against a full DB takes minutes (even
after the ZIP-snapshot speedup) — the browser's HTTP client times
out long before the work completes, leaving the user with a
'backend connection failed' modal and no way to know whether the
correction actually finished.

This module decouples the trigger from the work:

  POST /override/vulnrichment/start  → spawn thread, return job_id
  GET  /override/vulnrichment/status/{job_id} → progress + stats

State lives in a module-global dict so it survives the request but
not a backend restart. That's acceptable: a restart aborts any
in-flight job anyway, and the user just re-clicks.

Thread safety: dict assignment is atomic in CPython, and only the
worker thread mutates a given job's fields after creation. The
status reader returns whatever it sees — slightly stale reads are
fine for a 2-second poll loop.
"""
from __future__ import annotations

import logging
import threading
import uuid
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

from sqlalchemy.orm import Session

from app.database import SessionLocal

logger = logging.getLogger(__name__)


# job_id (str) → state dict
_jobs: Dict[str, Dict[str, Any]] = {}
# Cap the in-memory ring so a misbehaving caller can't OOM us.
_MAX_JOBS = 50


def _prune_oldest() -> None:
    if len(_jobs) <= _MAX_JOBS:
        return
    # Drop the oldest finished job. Finished = has finished_at.
    finished = [
        (jid, j.get("finished_at"))
        for jid, j in _jobs.items()
        if j.get("finished_at")
    ]
    if not finished:
        return
    finished.sort(key=lambda t: t[1])
    _jobs.pop(finished[0][0], None)


def start_vulnrichment_job(
    user_id: int,
    cve_ids: Optional[List[str]] = None,
    asset_ids: Optional[List[int]] = None,
    dry_run: bool = False,
) -> str:
    """Spawn a worker thread, return its job_id immediately."""
    job_id = str(uuid.uuid4())
    _jobs[job_id] = {
        "job_id": job_id,
        "type": "vulnrichment",
        "status": "queued",       # queued | running | completed | failed
        "user_id": user_id,
        "dry_run": dry_run,
        "started_at": datetime.now().isoformat(),
        "finished_at": None,
        "stage": "starting",      # human-readable progress label
        "total": 0,
        "done": 0,
        "updated": 0,
        "checked": 0,
        "not_found": 0,
        "error": None,
        "result": None,
    }
    _prune_oldest()

    thread = threading.Thread(
        target=_run_vulnrichment_job,
        args=(job_id, cve_ids, asset_ids, dry_run),
        daemon=True,
        name=f"override-{job_id[:8]}",
    )
    thread.start()
    return job_id


def get_job(job_id: str) -> Optional[Dict[str, Any]]:
    """Return a snapshot of the job state, or None if unknown."""
    job = _jobs.get(job_id)
    if not job:
        return None
    # Return a copy so the caller can't mutate the live state.
    return dict(job)


def list_jobs(limit: int = 10) -> List[Dict[str, Any]]:
    """Most-recent-first list of jobs (for an admin overview)."""
    snapshot = list(_jobs.values())
    snapshot.sort(key=lambda j: j.get("started_at") or "", reverse=True)
    return [dict(j) for j in snapshot[:limit]]


def _run_vulnrichment_job(
    job_id: str,
    cve_ids: Optional[List[str]],
    asset_ids: Optional[List[int]],
    dry_run: bool,
) -> None:
    """Worker thread body. Owns its own DB session."""
    # Import lazily to avoid a circular import at module load.
    from app.services.vuln_override_service import (
        VulnOverrideService,
        correct_vulnerability_scores,
    )
    from app.models.vulnerability import Vulnerability, VulnerabilityStatus

    job = _jobs[job_id]
    job["status"] = "running"
    db: Session = SessionLocal()
    try:
        # Stage 1: figure out target CVE set so progress has a known total
        job["stage"] = "selecting target CVEs"
        query = db.query(Vulnerability).filter(
            Vulnerability.status == VulnerabilityStatus.open
        )
        if cve_ids:
            query = query.filter(Vulnerability.cve_id.in_(cve_ids))
        if asset_ids:
            query = query.filter(Vulnerability.asset_id.in_(asset_ids))
        vulns = query.all()
        unique = sorted({
            v.cve_id for v in vulns
            if v.cve_id and not v.cve_id.startswith("NESSUS-")
        })
        job["total"] = len(unique)

        if not unique:
            job["status"] = "completed"
            job["stage"] = "no CVEs to correct"
            job["finished_at"] = datetime.now().isoformat()
            job["result"] = {"message": "no real CVEs in scope", "updated": 0}
            return

        # Stage 2: fetch (ZIP snapshot for large batches, else per-CVE raw)
        if len(unique) > VulnOverrideService._ZIP_FALLBACK_THRESHOLD:
            job["stage"] = (
                f"downloading vulnrichment ZIP snapshot (~249 MB) for {len(unique)} CVEs"
            )
        else:
            job["stage"] = f"fetching {len(unique)} CVEs from github raw"

        # Stage 3: apply overrides via existing helper (does the heavy lift)
        # We don't get per-CVE callbacks from inside the helper, so we
        # surface "done" only after it returns. The two stages above
        # already tell the user what's currently happening.
        result = correct_vulnerability_scores(
            db,
            cve_ids=unique,
            asset_ids=asset_ids,
            dry_run=dry_run,
        )
        job["done"] = job["total"]
        job["checked"] = result.get("checked", 0)
        job["updated"] = result.get("updated", 0)
        job["not_found"] = result.get("not_found", 0)
        job["result"] = result
        job["status"] = "completed"
        job["stage"] = (
            f"done — {job['updated']} updated, "
            f"{job['not_found']} not in vulnrichment feed"
        )
    except Exception as e:
        logger.exception("override job %s failed", job_id)
        job["status"] = "failed"
        job["error"] = str(e)
        job["stage"] = f"failed: {e}"
    finally:
        job["finished_at"] = datetime.now().isoformat()
        db.close()


# ----------------------------------------------------------------------
# Nessus scan → poll → auto-import job (Plan F)
# ----------------------------------------------------------------------

# How long to wait for a launched scan to finish before giving up.
_NESSUS_POLL_INTERVAL_SEC = 15
_NESSUS_POLL_TIMEOUT_SEC = 3600   # 1h — large network scans can run long


def start_nessus_scan_import_job(
    user_id: int,
    asset_id: Optional[int] = None,
    scan_id: Optional[int] = None,
    alt_targets: Optional[List[str]] = None,
) -> str:
    """Spawn a worker: launch a Nessus scan, poll until done, then import.

    asset_id optional — when set, the asset's IP is the alt_target and
    only that scan's results for that host get re-synced. When omitted,
    the configured scan runs against its own scope and the full sync
    imports everything.
    """
    job_id = str(uuid.uuid4())
    _jobs[job_id] = {
        "job_id": job_id,
        "type": "nessus_scan_import",
        "status": "queued",
        "user_id": user_id,
        "asset_id": asset_id,
        "scan_id": scan_id,
        "started_at": datetime.now().isoformat(),
        "finished_at": None,
        "stage": "starting",
        "nessus_status": None,
        "polls": 0,
        "error": None,
        "result": None,
    }
    _prune_oldest()
    thread = threading.Thread(
        target=_run_nessus_scan_import_job,
        args=(job_id, asset_id, scan_id, alt_targets),
        daemon=True,
        name=f"nessus-scan-{job_id[:8]}",
    )
    thread.start()
    return job_id


def _run_nessus_scan_import_job(
    job_id: str,
    asset_id: Optional[int],
    scan_id: Optional[int],
    alt_targets: Optional[List[str]],
) -> None:
    """Worker: launch scan → poll status → run_nessus_sync on completion."""
    import time as _time

    from app.integrations.nessus_client import (
        NessusClient, NessusAPIError, NessusAuthenticationError,
    )
    from app.services.nessus_sync import load_nessus_config, run_nessus_sync
    from app.models.asset import Asset

    job = _jobs[job_id]
    job["status"] = "running"
    db: Session = SessionLocal()
    try:
        config = load_nessus_config(db)
        if not config:
            raise RuntimeError("Nessus is not configured (Settings → Tenable Nessus)")

        targets = alt_targets
        if asset_id and not targets:
            asset = db.query(Asset).filter(Asset.id == asset_id).first()
            if not asset:
                raise RuntimeError(f"Asset {asset_id} not found")
            if not asset.ip_address:
                raise RuntimeError(f"Asset '{asset.hostname}' has no IP — cannot target")
            targets = [asset.ip_address]

        with NessusClient(
            base_url=config["base_url"],
            access_key=config["access_key"],
            secret_key=config["secret_key"],
            verify_ssl=bool(config.get("verify_ssl", True)),
        ) as client:
            # Resolve scan_id if not given.
            resolved_scan_id = scan_id
            if not resolved_scan_id:
                defaults = config.get("default_scan_ids") or []
                if defaults:
                    resolved_scan_id = int(defaults[0])
                else:
                    available = client.list_scans()
                    if not available:
                        raise RuntimeError("No Nessus scans visible for these API keys")
                    resolved_scan_id = int(available[0]["id"])
            job["scan_id"] = resolved_scan_id

            job["stage"] = f"launching scan {resolved_scan_id}"
            skip_poll = False
            status = "pending"
            from app.integrations.nessus_client import NessusAPIError as _NErr
            try:
                client.launch_scan(scan_id=resolved_scan_id, alt_targets=targets)
            except _NErr as e:
                # Nessus Essentials (free) often refuses API-triggered
                # launches — drops the connection, scan never enters
                # 'running'. If completed results already exist, import
                # them (flagged); else fail with an actionable message.
                job["launch_warning"] = str(e)
                pre_status = ""
                try:
                    pre_status = client.get_scan_status(resolved_scan_id)
                except Exception:
                    pass
                if pre_status in ("completed", "imported"):
                    job["stage"] = (
                        f"launch not accepted (Nessus Essentials?) — importing "
                        f"existing '{pre_status}' results for scan {resolved_scan_id}"
                    )
                    job["nessus_status"] = pre_status
                    status = pre_status
                    skip_poll = True
                else:
                    raise RuntimeError(
                        f"Could not launch scan {resolved_scan_id} via the Nessus API "
                        f"(no existing results to import). This Nessus edition "
                        f"(likely Essentials/free) may not allow API-triggered scans — "
                        f"launch it in the Nessus UI, then use 'Sync now'. ({e})"
                    )

            # Poll until terminal status (unless we're importing existing).
            if not skip_poll:
                deadline = _time.monotonic() + _NESSUS_POLL_TIMEOUT_SEC
                terminal = {"completed", "imported", "canceled", "aborted", "empty"}
                while _time.monotonic() < deadline:
                    _time.sleep(_NESSUS_POLL_INTERVAL_SEC)
                    job["polls"] += 1
                    try:
                        status = client.get_scan_status(resolved_scan_id)
                    except Exception as e:
                        logger.warning("nessus poll %s status error: %s", resolved_scan_id, e)
                        continue
                    job["nessus_status"] = status
                    job["stage"] = f"scan {resolved_scan_id} status: {status} (poll #{job['polls']})"
                    if status in terminal:
                        break
                else:
                    raise RuntimeError(
                        f"Scan {resolved_scan_id} did not finish within "
                        f"{_NESSUS_POLL_TIMEOUT_SEC // 60} min (last status: {status})"
                    )

                if status in ("canceled", "aborted"):
                    raise RuntimeError(f"Scan {resolved_scan_id} ended as '{status}' — no import")

        # Scan done → import results (own session inside run_nessus_sync).
        job["stage"] = f"importing results from scan {resolved_scan_id}"
        sync_stats = run_nessus_sync(db, scan_ids=[resolved_scan_id])
        db.commit()
        job["result"] = sync_stats
        job["status"] = "completed"
        job["stage"] = (
            f"done — {sync_stats.get('vulns_created', 0)} created, "
            f"{sync_stats.get('vulns_merged', 0)} merged, "
            f"{sync_stats.get('vulns_marked_patched', 0)} patched"
        )
    except (NessusAPIError, NessusAuthenticationError) as e:
        logger.exception("nessus scan-import job %s API error", job_id)
        job["status"] = "failed"
        job["error"] = str(e)
        job["stage"] = f"nessus API error: {e}"
    except Exception as e:
        logger.exception("nessus scan-import job %s failed", job_id)
        job["status"] = "failed"
        job["error"] = str(e)
        job["stage"] = f"failed: {e}"
    finally:
        job["finished_at"] = datetime.now().isoformat()
        db.close()