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 typing import Optional

import httpx

from config import settings
from docker_client import DockerError, get_client, safe_call
from models.setting import EVENT_UPDATE_AVAILABLE
from services import notify_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

    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()


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


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)
    else:
        # find last colon after last slash
        slash = ref.rfind("/")
        colon = ref.rfind(":")
        if colon > slash:
            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]
    else:
        registry = "registry-1.docker.io"
        repo = ref
        if "/" not in repo:
            repo = "library/" + repo
    return registry, repo, tag


def _local_digest(image: str) -> Optional[str]:
    try:
        client = get_client()
        img = safe_call(client.images.get, image)
    except DockerError:
        return None
    repo_digests = img.attrs.get("RepoDigests") or []
    for rd in repo_digests:
        if "@" in rd:
            return rd.split("@", 1)[1]
    return None


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


async def _get_token(client: httpx.AsyncClient, www_auth: str) -> Optional[str]:
    # 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, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        return data.get("token") or data.get("access_token")
    except (httpx.HTTPError, ValueError):
        return None


async def remote_digest(image: str) -> Optional[str]:
    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}
    async with httpx.AsyncClient(follow_redirects=True) as client:
        try:
            resp = await client.head(url, headers=headers, timeout=10)
            if resp.status_code == 401:
                token = await _get_token(client, resp.headers.get("WWW-Authenticate", ""))
                if not token:
                    return None
                headers["Authorization"] = f"Bearer {token}"
                resp = await client.head(url, headers=headers, timeout=10)
            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, timeout=10)
            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:
    local = _local_digest(image)
    remote = await remote_digest(image)
    error = None
    if remote 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,
    )
    _CACHE[image] = status
    if update_available and image not in _NOTIFIED:
        _NOTIFIED.add(image)
        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)
    elif not update_available:
        _NOTIFIED.discard(image)
    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


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


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)
        # Re-read the interval each cycle so Settings changes take effect.
        interval = max(settings_service.get_update_interval(), 5) * 60
        await asyncio.sleep(interval)