stackpilot/backend/routers/images.py

"""Image listing, update checks and vulnerability scanning."""
from __future__ import annotations

import json

from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
from sqlmodel import Session

from auth import get_current_user, require_admin
from database import engine, get_session
from models.image_scan import ImageScan, ScanDetail, ScanRequest, ScanSummary
from models.user import User
from services import audit_service, image_service, scan_service, update_service

router = APIRouter(prefix="/api/images", tags=["images"])


def _ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


@router.get("")
def list_images(_user: User = Depends(get_current_user)) -> list[dict]:
    return image_service.list_images()


@router.get("/updates")
def updates(_user: User = Depends(get_current_user)) -> dict:
    return update_service.get_cache()


@router.post("/check")
async def check(_user: User = Depends(require_admin)) -> dict:
    return await update_service.check_all()


@router.post("/prune")
def prune(
    request: Request,
    response: Response,
    all_unused: bool = Query(False, alias="all"),
    background: bool = Query(False),
    user: User = Depends(require_admin),
) -> dict:
    """Remove unused images (``all=true``: also unused tagged ones, ``-a``).

    Waits for the result by default. ``background=true`` returns 202 with the
    initial status at once and leaves the caller to poll /prune-status, which
    is what the UI does: a large prune can outlast the proxy's read timeout.
    """
    username, ip = user.username, _ip(request)

    def audit(result: dict) -> None:
        # Docker lists every untagging too; only "Deleted" entries are images.
        deleted = sum(1 for entry in result["ImagesDeleted"] if "Deleted" in entry)
        # Its own session: in the background case the request's is long closed.
        with Session(engine) as own:
            audit_service.record(
                own, user=username, action="image.prune", target="*",
                detail=(
                    f"all={all_unused} deleted={deleted} "
                    f"reclaimed={result['SpaceReclaimed']}"
                ),
                ip=ip,
            )

    try:
        if background:
            response.status_code = 202
            return image_service.start_prune(all_unused, on_done=audit)
        result = image_service.prune_images(all_unused)
    except image_service.PruneBusy as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from exc
    audit(result)
    return result


@router.get("/prune-status")
def prune_status(_admin: User = Depends(require_admin)) -> dict:
    """Progress of a running (or the last) prune. Admin-only like the prune."""
    return image_service.prune_status()


# --------------------------------------------------------------------------- #
# Vulnerability scanning
# --------------------------------------------------------------------------- #


def _summary(row: ImageScan, stale: bool = False) -> ScanSummary:
    return ScanSummary(
        image=row.image,
        digest=row.digest,
        scanner=row.scanner,
        critical=row.critical,
        high=row.high,
        medium=row.medium,
        low=row.low,
        unknown=row.unknown,
        fixable=row.fixable,
        total=scan_service.total(row),
        scanned_at=row.scanned_at,
        duration_ms=row.duration_ms,
        error=row.error,
        stale=stale,
    )


@router.get("/scans", response_model=list[ScanSummary])
def list_scans(
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> list[ScanSummary]:
    """Every cached scan result. Cheap — no scanning happens here."""
    rows = scan_service.all_scans(session).values()
    return [_summary(row, scan_service.is_stale(row)) for row in rows]


@router.get("/scan", response_model=ScanDetail)
def scan_detail(
    image: str = Query(...),
    session: Session = Depends(get_session),
    _user: User = Depends(get_current_user),
) -> ScanDetail:
    row = scan_service.all_scans(session).get(image)
    if not row:
        raise HTTPException(status_code=404, detail=f"'{image}' has not been scanned")
    try:
        findings = json.loads(row.findings or "[]")
    except json.JSONDecodeError:
        findings = []
    return ScanDetail(
        **_summary(row, scan_service.is_stale(row)).model_dump(), findings=findings
    )


@router.post("/scan", response_model=ScanSummary)
async def scan_image(
    body: ScanRequest,
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> ScanSummary:
    """Scan one image. Takes a while — the scanner runs as a container."""
    image = (body.image or "").strip()
    if not image:
        raise HTTPException(status_code=400, detail="An image is required")
    row = await scan_service.scan(session, image)
    audit_service.record(
        session, user=user.username, action="image.scan", target=image,
        detail=row.error or f"{scan_service.total(row)} finding(s), {row.fixable} fixable",
        ip=_ip(request),
    )
    if row.error:
        # The row is stored either way, so the UI can show why it failed.
        raise HTTPException(status_code=502, detail=row.error)
    return _summary(row)


@router.post("/scan-all")
async def scan_all(
    request: Request,
    session: Session = Depends(get_session),
    user: User = Depends(require_admin),
) -> dict:
    """Scan every image a stack is using. Serial, and one sweep at a time."""
    images = sorted({row["tag"] for row in image_service.list_images() if row["stacks"]})
    try:
        result = await scan_service.sweep(session, images)
    except scan_service.ScanError as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from exc
    audit_service.record(
        session, user=user.username, action="image.scan-all", target="*",
        detail=f"{result['scanned']} scanned, {result['failed']} failed",
        ip=_ip(request),
    )
    return result


@router.get("/scan-status")
def scan_status(_user: User = Depends(get_current_user)) -> dict:
    """Progress of a running sweep, so the button can show it."""
    return scan_service.sweep_status()