stackpilot/backend/services/update_service.py

"""Image update checker.

Compares the locally-pulled manifest digest (from RepoDigests) against the
current manifest digest in the registry. Supports Docker Hub, ghcr.io, lscr.io
and other token-auth v2 registries.
"""
from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import asdict, dataclass
from collections.abc import Callable
from typing import Optional

import httpx

from docker_client import DockerError, get_client, safe_call
from models.setting import EVENT_UPDATE_AVAILABLE
from services import notify_service, registry_service, settings_service

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

_MANIFEST_ACCEPT = ", ".join(
    [
        "application/vnd.docker.distribution.manifest.v2+json",
        "application/vnd.docker.distribution.manifest.list.v2+json",
        "application/vnd.oci.image.manifest.v1+json",
        "application/vnd.oci.image.index.v1+json",
    ]
)


@dataclass
class UpdateStatus:
    image: str
    update_available: bool
    current_digest: Optional[str]
    remote_digest: Optional[str]
    checked_at: float
    error: Optional[str] = None
    #: Built on this host (``docker build``, compose ``build:``) rather than
    #: pulled: there is no registry copy to compare against, so it is never
    #: "out of date" — and asking Docker Hub about it only got a 401 that read
    #: as "needs credentials".
    local: bool = False

    def to_dict(self) -> dict:
        return asdict(self)


# image ref -> UpdateStatus
_CACHE: dict[str, UpdateStatus] = {}

# images we've already sent an "update available" notification for, so the
# background loop doesn't re-notify on every cycle.
_NOTIFIED: set[str] = set()

#: Optional sink for cache writes.
#:
#: This module is pure registry logic and knows nothing about storage, which
#: keeps it unit-testable without a database. ``main.lifespan`` registers a
#: callback that mirrors each entry into SQLite (see
#: ``services/image_status_store.py``) and seeds the cache from it at startup.
#: Without it a restart blanked every update badge until the next background
#: sweep — up to an hour — and re-announced updates it had already notified
#: about.
_persist_cb: Optional[Callable[[UpdateStatus, bool], None]] = None


#: Optional sink for "these images are still in use", same opt-in shape as
#: _persist_cb. Keeps both the dict and the table from growing one entry per
#: image tag that was ever running, for the life of the install.
_prune_cb: Optional[Callable[[set], int]] = None


def set_persist_callback(
    callback: Optional[Callable[[UpdateStatus, bool], None]],
    prune: Optional[Callable[[set], int]] = None,
) -> None:
    global _persist_cb, _prune_cb
    _persist_cb = callback
    _prune_cb = prune


def restore_cache(entries: list[tuple[UpdateStatus, bool]]) -> None:
    """Seed the in-memory cache from persisted rows at startup."""
    for status, notified in entries:
        _CACHE[status.image] = status
        if notified:
            _NOTIFIED.add(status.image)


def _store(status: UpdateStatus, notified: bool) -> None:
    _CACHE[status.image] = status
    if _persist_cb is not None:
        try:
            _persist_cb(status, notified)
        except Exception as exc:  # noqa: BLE001 - persistence is best-effort
            logger.debug("Could not persist update status for %s: %s", status.image, exc)


# --------------------------------------------------------------------------- #
# Image reference parsing
# --------------------------------------------------------------------------- #


#: Docker Hub under the names people write in compose files. Only
#: registry-1.docker.io serves the v2 API; https://docker.io/v2/ redirects to
#: the marketing site, so a `docker.io/library/postgres` image could never be
#: checked and sat at "could not reach registry" forever.
_DOCKER_HUB = "registry-1.docker.io"
_DOCKER_HUB_ALIASES = {"docker.io", "index.docker.io", "registry.docker.io", _DOCKER_HUB}


def parse_ref(image: str) -> tuple[str, str, str]:
    """Return (registry_host, repository, tag/digest)."""
    ref = image
    tag = "latest"
    # split tag (but not the registry port colon)
    if "@" in ref:
        ref, tag = ref.split("@", 1)
    # find last colon after last slash. With a digest present this is the tag
    # in "name:tag@sha256:…", which the digest outranks — but it must not stay
    # glued to the repository name.
    slash = ref.rfind("/")
    colon = ref.rfind(":")
    if colon > slash:
        if "@" not in image:
            tag = ref[colon + 1 :]
        ref = ref[:colon]

    parts = ref.split("/", 1)
    if len(parts) == 2 and ("." in parts[0] or ":" in parts[0] or parts[0] == "localhost"):
        registry = parts[0]
        repo = parts[1]
        if registry.lower() in _DOCKER_HUB_ALIASES:
            registry = _DOCKER_HUB
    else:
        registry = _DOCKER_HUB
        repo = ref
    if registry == _DOCKER_HUB and "/" not in repo:
        repo = "library/" + repo
    return registry, repo, tag


def _local_attrs(image: str) -> Optional[dict]:
    try:
        return safe_call(get_client().images.get, image).attrs
    except DockerError:
        return None


def _built_locally(attrs: Optional[dict]) -> bool:
    """Whether the image came from a build on this host, not from a registry.

    Docker 29+ records where an image came from (``Identity.Pull`` /
    ``Identity.Build``). Older daemons do not; there a built image is the one
    without RepoDigests. Unknown means "pulled" — skipping the check for a
    pulled image would hide its updates for good.
    """
    if attrs is None:
        return False
    identity = attrs.get("Identity")
    if isinstance(identity, dict) and identity:
        return bool(identity.get("Build")) and not identity.get("Pull")
    return not (attrs.get("RepoDigests") or [])


def _local_digest(image: str, attrs: Optional[dict] = None) -> Optional[str]:
    if attrs is None:
        attrs = _local_attrs(image)
    if attrs is None:
        return None
    repo_digests = [rd for rd in attrs.get("RepoDigests") or [] if "@" in rd]
    # An image can carry digests for several repositories (the same image
    # pulled as lscr.io/... and as docker.io/...); compare against the one
    # this reference is checked against, not whichever comes first.
    wanted = parse_ref(image)[:2]
    for rd in repo_digests:
        name, digest = rd.split("@", 1)
        if parse_ref(name)[:2] == wanted:
            return digest
    return repo_digests[0].split("@", 1)[1] if repo_digests else None


# --------------------------------------------------------------------------- #
# Registry manifest digest
# --------------------------------------------------------------------------- #


async def _get_token(
    client: httpx.AsyncClient,
    www_auth: str,
    auth: Optional[tuple[str, str]] = None,
) -> Optional[str]:
    """Follow a Bearer challenge, with credentials when we have them.

    A public image gets an anonymous token; a private one only gets a token at
    all if the request to the token realm is authenticated.
    """
    # Parse: Bearer realm="...",service="...",scope="..."
    params = {}
    if not www_auth.lower().startswith("bearer"):
        return None
    for part in www_auth[len("Bearer ") :].split(","):
        if "=" in part:
            k, v = part.split("=", 1)
            params[k.strip()] = v.strip().strip('"')
    realm = params.pop("realm", None)
    if not realm:
        return None
    try:
        resp = await client.get(realm, params=params, auth=auth, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        return data.get("token") or data.get("access_token")
    except (httpx.HTTPError, ValueError):
        return None


class AuthRequired(Exception):
    """The registry wants credentials we do not have (or rejected ours).

    Distinct from "could not reach the registry" on purpose: a private image
    with no configured credentials used to be indistinguishable from a network
    blip, so the UI said nothing and the stack looked up to date forever.
    """


async def remote_digest(image: str) -> Optional[str]:
    """The digest the registry currently serves for this tag.

    Raises :class:`AuthRequired` when the registry refuses us; returns None when
    it could not be reached or answered without a digest.
    """
    registry, repo, tag = parse_ref(image)
    if tag.startswith("sha256:"):
        return tag
    scheme = "https"
    url = f"{scheme}://{registry}/v2/{repo}/manifests/{tag}"
    headers = {"Accept": _MANIFEST_ACCEPT}
    auth = registry_service.credentials_for(registry)
    async with httpx.AsyncClient(follow_redirects=True) as client:
        try:
            resp = await client.head(url, headers=headers, timeout=10)
            if resp.status_code == 401:
                challenge = resp.headers.get("WWW-Authenticate", "")
                if challenge.lower().startswith("basic"):
                    # A plain htpasswd-protected registry: no token dance.
                    if not auth:
                        raise AuthRequired(registry)
                    resp = await client.head(url, headers=headers, auth=auth, timeout=10)
                else:
                    token = await _get_token(client, challenge, auth)
                    if not token:
                        raise AuthRequired(registry)
                    headers["Authorization"] = f"Bearer {token}"
                    resp = await client.head(url, headers=headers, timeout=10)
            if resp.status_code in (401, 403):
                raise AuthRequired(registry)
            if resp.status_code == 405 or "Docker-Content-Digest" not in resp.headers:
                # Some registries don't support HEAD; fall back to GET.
                resp = await client.get(
                    url, headers=headers, auth=auth if "Authorization" not in headers else None,
                    timeout=10,
                )
                if resp.status_code in (401, 403):
                    raise AuthRequired(registry)
            digest = resp.headers.get("Docker-Content-Digest")
            return digest
        except httpx.HTTPError as exc:
            logger.debug("remote_digest failed for %s: %s", image, exc)
            return None


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #


async def check_image(image: str) -> UpdateStatus:
    attrs = _local_attrs(image)
    if _built_locally(attrs):
        status = UpdateStatus(
            image=image,
            update_available=False,
            current_digest=None,
            remote_digest=None,
            checked_at=time.time(),
            local=True,
        )
        _NOTIFIED.discard(image)
        _store(status, False)
        return status
    local = _local_digest(image, attrs)
    error = None
    try:
        remote = await remote_digest(image)
    except AuthRequired as exc:
        # Say which registry, because the fix is to add credentials for it.
        remote = None
        error = (
            f"{exc} needs credentials"
            if not registry_service.credentials_for(str(exc))
            else f"{exc} rejected the stored credentials"
        )
    if remote is None and error is None:
        error = "could not reach registry"
    update_available = bool(local and remote and local != remote)
    status = UpdateStatus(
        image=image,
        update_available=update_available,
        current_digest=local,
        remote_digest=remote,
        checked_at=time.time(),
        error=error,
    )
    if update_available and image not in _NOTIFIED:
        # Marked before the attempt, not after: a notifier that is down should
        # not make every cycle re-announce the same update.
        _NOTIFIED.add(image)
        _store(status, True)
        try:
            await notify_service.notify(
                EVENT_UPDATE_AVAILABLE,
                "Image update available",
                f"A newer image is available for {image}.",
            )
        except Exception as exc:  # noqa: BLE001 - notifications are best-effort
            logger.debug("update notify failed for %s: %s", image, exc)
    else:
        if not update_available:
            _NOTIFIED.discard(image)
        _store(status, image in _NOTIFIED)
    return status


def _all_running_images() -> set[str]:
    images: set[str] = set()
    try:
        client = get_client()
        for c in safe_call(client.containers.list, all=True):
            cfg_image = c.attrs.get("Config", {}).get("Image")
            if cfg_image:
                images.add(cfg_image)
    except DockerError:
        pass
    return images


def stack_images(stack_id: str) -> set[str]:
    """Images used by the containers of one compose project (== stack id)."""
    from services.compose_service import project_name

    project = project_name(stack_id)
    images: set[str] = set()
    try:
        client = get_client()
        for c in safe_call(client.containers.list, all=True):
            labels = c.labels or {}
            if labels.get("com.docker.compose.project") != project:
                continue
            cfg_image = c.attrs.get("Config", {}).get("Image")
            if cfg_image:
                images.add(cfg_image)
    except DockerError:
        pass
    return images


def stacks_update_summary() -> dict[str, dict]:
    """Per-stack image-update status for every running compose project, read
    from the digest cache the background loop maintains — no registry calls, so
    it's cheap enough for the stacks list to poll. Stacks with no cached image
    yet are simply absent (treated as "no update" by the UI)."""
    by_stack: dict[str, set[str]] = {}
    try:
        client = get_client()
        for c in safe_call(client.containers.list, all=True):
            project = (c.labels or {}).get("com.docker.compose.project")
            if not project:
                continue
            cfg_image = c.attrs.get("Config", {}).get("Image")
            if cfg_image:
                by_stack.setdefault(project, set()).add(cfg_image)
    except DockerError:
        return {}

    summary: dict[str, dict] = {}
    for stack_id, images in by_stack.items():
        stale = [
            img
            for img in images
            if (st := _CACHE.get(img)) is not None and st.update_available
        ]
        summary[stack_id] = {
            "update_available": bool(stale),
            "stale_images": stale,
        }
    return summary


async def stack_updates(stack_id: str, refresh: bool = True) -> dict:
    """Update status for one stack's images.

    ``refresh=True`` queries the registry now; ``False`` reads the cache the
    background loop already populated (so the auto-update pass adds no extra
    registry round-trips).
    """
    images = stack_images(stack_id)
    result: dict[str, dict] = {}
    for image in images:
        status = await check_image(image) if refresh else _CACHE.get(image)
        if status is not None:
            result[image] = status.to_dict()
    stale = [img for img, st in result.items() if st.get("update_available")]
    return {
        "stack_id": stack_id,
        "update_available": bool(stale),
        "stale_images": stale,
        "images": result,
    }


def refresh_stack_local(stack_id: str) -> None:
    """Re-read the local digests of one stack's images and reconcile them with
    the cached remote digests (no registry calls). Called right after a manual
    pull/update so the amber indicator clears immediately instead of lingering
    until the next background pass."""
    for image in stack_images(stack_id):
        status = _CACHE.get(image)
        if status is None:
            continue
        local = _local_digest(image)
        status.current_digest = local
        status.update_available = bool(
            local and status.remote_digest and local != status.remote_digest
        )
        status.checked_at = time.time()
        if not status.update_available:
            _NOTIFIED.discard(image)
        # Persisted too, or a restart brought the cleared badge back until the
        # next background check.
        _store(status, image in _NOTIFIED)


async def check_all() -> dict[str, dict]:
    images = _all_running_images()
    for image in images:
        await check_image(image)
    _forget_unused(set(images))
    return {k: v.to_dict() for k, v in _CACHE.items()}


def _forget_unused(keep: set) -> None:
    """Drop images no running container references any more."""
    for image in [i for i in _CACHE if i not in keep]:
        del _CACHE[image]
        _NOTIFIED.discard(image)
    if _prune_cb is not None:
        try:
            _prune_cb(keep)
        except Exception as exc:  # noqa: BLE001 - housekeeping is best-effort
            logger.debug("Could not prune persisted update statuses: %s", exc)


def get_cache() -> dict[str, dict]:
    return {k: v.to_dict() for k, v in _CACHE.items()}


async def background_loop():
    # initial delay so startup isn't blocked
    await asyncio.sleep(30)
    while True:
        try:
            await check_all()
            logger.info("Image update check complete (%d images)", len(_CACHE))
        except Exception as exc:  # noqa: BLE001
            logger.warning("Image update check failed: %s", exc)
        # Apply auto-update policies using the digest cache we just refreshed.
        # Lazy import avoids a circular import (auto_update_service imports us).
        try:
            from services import auto_update_service

            await auto_update_service.run_due()
        except Exception as exc:  # noqa: BLE001
            logger.warning("Auto-update pass failed: %s", exc)
        # Re-read the interval each cycle so Settings changes take effect.
        interval = max(settings_service.get_update_interval(), 5) * 60
        await asyncio.sleep(interval)