stackpilot/backend/services/scan_service.py

"""Vulnerability scanning for the images stacks actually run.

Scanning runs Trivy in a throwaway container — the same helper-container trick
backups already use — rather than adding a scanner to StackPilot's own image.
That keeps a 100 MB security tool and its weekly-changing vulnerability database
out of the release, and means the scanner updates itself by pulling a newer tag.

The container gets the Docker socket read-only, so it inspects images the daemon
already has instead of pulling them again from a registry, and a named volume
for its database, so the ~50 MB download happens once rather than per scan.

Findings are stored worst-first and capped. A base image with 400 known CVEs is
common and the tail is unreadable; what matters is how many there are, how many
have a fix, and what the worst ones are.
"""
from __future__ import annotations

import asyncio
import json
import logging
import threading
import time
from datetime import datetime, timezone
from typing import Optional

from sqlmodel import Session, select

from config import settings
from docker_client import DockerError, get_client, get_long_client, run_container, safe_call
from models.image_scan import SEVERITIES, ImageScan

logger = logging.getLogger("stackpilot.scan")

#: Named volume for Trivy's vulnerability database.
CACHE_VOLUME = "stackpilot-trivy-cache"

#: Kept per image. Enough to see the shape of the problem, not enough to turn a
#: database row into a megabyte of JSON nobody reads.
MAX_FINDINGS = 200

SCAN_TIMEOUT = 900.0

#: One sweep at a time. Scanning is CPU- and IO-heavy, and two passes over the
#: same images would just fight each other for the daemon.
_sweep_lock = threading.Lock()
_sweeping = {"active": False, "done": 0, "total": 0, "scanned": 0, "failed": 0, "run": 0}


class ScanError(Exception):
    """The scanner could not run, or produced something unusable."""


# --------------------------------------------------------------------------- #
# Running the scanner
# --------------------------------------------------------------------------- #


#: Wait for the scanner a little longer than its own --timeout of 10m.
SCAN_WAIT_SECONDS = 720


def _run_trivy(image: str) -> str:
    """Run Trivy against a local image and return its JSON report."""
    # Trivy may take its full --timeout (the first run also downloads its
    # vulnerability database); the 60s default client gave up long before.
    client = get_long_client()
    socket_path = settings.DOCKER_SOCKET
    command = [
        "image",
        "--format", "json",
        "--quiet",
        # Vulnerabilities only: secret and misconfiguration scanning are a
        # different feature with a very different false-positive profile.
        "--scanners", "vuln",
        "--timeout", "10m",
        image,
    ]
    try:
        output = run_container(
            client,
            settings.SCANNER_IMAGE,
            command,
            volumes={
                socket_path: {"bind": "/var/run/docker.sock", "mode": "ro"},
                CACHE_VOLUME: {"bind": "/root/.cache/trivy", "mode": "rw"},
            },
            timeout=SCAN_WAIT_SECONDS,
        )
    except DockerError as exc:
        raise ScanError(_readable(Exception(exc.detail or exc.error))) from exc
    except Exception as exc:  # noqa: BLE001
        raise ScanError(_readable(exc)) from exc
    return (output or b"").decode("utf-8", "replace")


def _readable(exc: Exception) -> str:
    """A scanner failure a person can act on, not a docker-py repr."""
    text = str(exc)
    if "No such image" in text or "not found" in text.lower():
        return "The scanner image could not be pulled, or the image no longer exists locally"
    return text[:500]


def _parse(report: str) -> tuple[dict[str, int], int, list[dict]]:
    """Counts by severity, how many are fixable, and the findings themselves."""
    try:
        data = json.loads(report or "{}")
    except json.JSONDecodeError as exc:
        raise ScanError("The scanner returned output that is not JSON") from exc

    counts = {name: 0 for name in SEVERITIES}
    fixable = 0
    findings: list[dict] = []
    for result in data.get("Results") or []:
        target = result.get("Target") or ""
        for vuln in result.get("Vulnerabilities") or []:
            severity = str(vuln.get("Severity") or "unknown").lower()
            if severity not in counts:
                severity = "unknown"
            counts[severity] += 1
            fixed = vuln.get("FixedVersion") or ""
            if fixed:
                fixable += 1
            findings.append(
                {
                    "id": vuln.get("VulnerabilityID"),
                    "severity": severity,
                    "package": vuln.get("PkgName"),
                    "installed": vuln.get("InstalledVersion"),
                    "fixed": fixed or None,
                    "title": (vuln.get("Title") or "")[:300] or None,
                    "url": vuln.get("PrimaryURL"),
                    "target": target,
                }
            )

    findings.sort(key=lambda f: (SEVERITIES.index(f["severity"]), f["id"] or ""))
    return counts, fixable, findings[:MAX_FINDINGS]


def local_digest(image: str) -> Optional[str]:
    """The digest of the image the daemon currently has for this tag."""
    try:
        client = get_client()
        img = safe_call(client.images.get, image)
    except DockerError:
        return None
    return img.id  # the content-addressable local id, which changes on a pull


# --------------------------------------------------------------------------- #
# Scanning and storing
# --------------------------------------------------------------------------- #


async def scan(session: Session, image: str) -> ImageScan:
    """Scan one image and store the result, replacing any previous one."""
    started = time.monotonic()
    row = session.exec(select(ImageScan).where(ImageScan.image == image)).first()
    if row is None:
        row = ImageScan(image=image)

    try:
        # docker-py is synchronous and a scan takes minutes; off the event loop
        # it goes, or every other request waits behind it.
        report = await asyncio.wait_for(
            asyncio.to_thread(_run_trivy, image), timeout=SCAN_TIMEOUT
        )
        counts, fixable, findings = _parse(report)
    except asyncio.TimeoutError:
        return _store_error(session, row, image, "The scan timed out", started)
    except ScanError as exc:
        return _store_error(session, row, image, str(exc), started)

    row.digest = local_digest(image)
    row.scanner = "trivy"
    for name in SEVERITIES:
        setattr(row, name, counts[name])
    row.fixable = fixable
    row.findings = json.dumps(findings)
    row.scanned_at = datetime.now(timezone.utc)
    row.duration_ms = int((time.monotonic() - started) * 1000)
    row.error = None
    session.add(row)
    session.commit()
    session.refresh(row)
    return row


def _store_error(
    session: Session, row: ImageScan, image: str, message: str, started: float
) -> ImageScan:
    """Record a failed scan, keeping whatever counts the last good one found.

    A scanner that cannot run must not silently turn into "no vulnerabilities".
    """
    row.image = image
    row.error = message[:500]
    row.scanned_at = datetime.now(timezone.utc)
    row.duration_ms = int((time.monotonic() - started) * 1000)
    session.add(row)
    session.commit()
    session.refresh(row)
    return row


def total(row: ImageScan) -> int:
    return sum(getattr(row, name) for name in SEVERITIES)


def is_stale(row: ImageScan) -> bool:
    """Has the image been pulled since it was scanned?"""
    if not row.digest or row.error:
        return False
    current = local_digest(row.image)
    return bool(current and current != row.digest)


def all_scans(session: Session) -> dict[str, ImageScan]:
    return {row.image: row for row in session.exec(select(ImageScan)).all()}


def forget(session: Session, images: set[str]) -> int:
    """Drop scans for images that no longer exist. Returns how many went."""
    removed = 0
    for row in session.exec(select(ImageScan)).all():
        if row.image not in images:
            session.delete(row)
            removed += 1
    if removed:
        session.commit()
    return removed


# --------------------------------------------------------------------------- #
# Scanning everything
# --------------------------------------------------------------------------- #


def sweep_status() -> dict:
    return dict(_sweeping)


def _begin(total: int) -> None:
    if not _sweep_lock.acquire(blocking=False):
        raise ScanError("A scan is already running")
    _sweeping.update(
        {"active": True, "done": 0, "total": total, "scanned": 0, "failed": 0,
         "run": _sweeping.get("run", 0) + 1}
    )


async def _run(session: Session, images: list[str]) -> dict:
    failed = 0
    try:
        for image in images:
            try:
                row = await scan(session, image)
                if row.error:
                    failed += 1
            except Exception as exc:  # noqa: BLE001 - one image must not stop the sweep
                logger.warning("Scan failed for %s: %s", image, exc)
                failed += 1
            _sweeping["done"] += 1
    finally:
        _sweeping.update({"active": False, "scanned": len(images), "failed": failed})
        _sweep_lock.release()
    return {"scanned": len(images), "failed": failed}


async def sweep(session: Session, images: list[str]) -> dict:
    """Scan a list of images one after another.

    Deliberately serial: each scan is already CPU-hungry, and running several at
    once on a homelab box would starve the very containers being protected.
    """
    _begin(len(images))
    return await _run(session, images)


#: Background sweeps, kept referenced until they finish.
_tasks: set[asyncio.Task] = set()


def start_sweep(images: list[str], on_done=None) -> dict:
    """Start a sweep in the background and return its status at once.

    A sweep over a dozen images takes far longer than the proxy waits for a
    response (600s), so the button used to report a 504 while the scans went
    on. Marked active before this returns, so a poll right after sees it.
    """
    from database import engine

    _begin(len(images))

    async def work() -> None:
        with Session(engine) as session:
            result = await _run(session, images)
            if on_done is not None:
                try:
                    on_done(session, result)
                except Exception:  # noqa: BLE001 - the sweep itself is done
                    logger.exception("Post-sweep hook failed")

    task = asyncio.create_task(work(), name="image-scan-sweep")
    _tasks.add(task)
    task.add_done_callback(_tasks.discard)
    return sweep_status()