stackpilot/backend/services/registry_service.py

"""Credentials for private container registries.

Two very different consumers need these, which is why this module exists rather
than the credentials living next to either of them:

* **StackPilot's own update checker** (``services/update_service.py``) talks to
  the registry v2 API over HTTP itself. Without credentials a private repository
  answers 401, the check gave up, and the UI said nothing at all — a stack could
  sit on a stale image for months and look up to date. That was the actual bug
  here, not a missing feature.
* **The Docker CLI**, which runs ``docker compose pull``. It reads its own
  ``config.json``, so this module writes one into ``${DATA_DIR}/docker`` and
  compose runs with ``DOCKER_CONFIG`` pointed at it.

Lookups happen inside async registry calls that have no database session, so
the rows are mirrored into a small in-memory cache. :func:`reload` refills it and
rewrites the CLI config; every write path calls it, and so does startup.

Passwords are encrypted at rest and only ever decrypted into this cache and the
CLI config file (0600, in the data volume). They are never returned by the API.
"""
from __future__ import annotations

import base64
import json
import logging
import os
import threading
from typing import Optional

import httpx
from sqlmodel import Session, select

from config import settings
from models.registry import Registry
from services import crypto_service

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

#: What `parse_ref` calls Docker Hub, and what the CLI calls it. Every spelling
#: users type — docker.io, index.docker.io, the v1 URL — normalizes to the
#: first; the second is what has to appear in config.json for `docker pull`.
DOCKER_HUB = "registry-1.docker.io"
DOCKER_HUB_CONFIG_KEY = "https://index.docker.io/v1/"
_DOCKER_HUB_ALIASES = {
    "docker.io",
    "index.docker.io",
    "registry.docker.io",
    "registry-1.docker.io",
    "https://index.docker.io/v1/",
    "index.docker.io/v1/",
}

TIMEOUT = httpx.Timeout(15.0)

_lock = threading.Lock()
_credentials: dict[str, tuple[str, str]] = {}


class RegistryError(Exception):
    """A registry that cannot be reached, or credentials it rejects."""


# --------------------------------------------------------------------------- #
# Host normalization
# --------------------------------------------------------------------------- #


def canonical_host(value: str) -> str:
    """The registry host as an image reference would name it.

    Accepts what people actually paste: a bare host, a URL with a scheme, a
    trailing slash, Docker Hub under any of its names. Without this, credentials
    entered as ``docker.io`` would never be found for an image that parses as
    ``registry-1.docker.io``.
    """
    host = (value or "").strip().lower()
    if not host:
        raise RegistryError("Registry host is required")
    if host in _DOCKER_HUB_ALIASES:
        return DOCKER_HUB
    for scheme in ("https://", "http://"):
        if host.startswith(scheme):
            host = host[len(scheme) :]
            break
    host = host.split("/", 1)[0].rstrip("/")
    if host in _DOCKER_HUB_ALIASES:
        return DOCKER_HUB
    if not host:
        raise RegistryError("Registry host is required")
    return host


def is_docker_hub(host: str) -> bool:
    return canonical_host(host) == DOCKER_HUB


# --------------------------------------------------------------------------- #
# The cache the async paths read
# --------------------------------------------------------------------------- #


def reload(session: Session) -> int:
    """Refill the credential cache from the database and rewrite config.json.

    Called at startup and after every write. Returns how many registries are
    configured.
    """
    fresh: dict[str, tuple[str, str]] = {}
    for row in session.exec(select(Registry)).all():
        try:
            password = crypto_service.decrypt(row.password)
        except crypto_service.DecryptError as exc:
            # One unreadable row must not take the others down with it.
            logger.warning("Ignoring credentials for %s: %s", row.host, exc)
            continue
        fresh[row.host] = (row.username, password)
    with _lock:
        _credentials.clear()
        _credentials.update(fresh)
    _write_docker_config(fresh)
    return len(fresh)


def credentials_for(host: str) -> Optional[tuple[str, str]]:
    """(username, password) for a registry host, or None."""
    try:
        key = canonical_host(host)
    except RegistryError:
        return None
    with _lock:
        return _credentials.get(key)


def credentials_for_image(image: str) -> Optional[tuple[str, str]]:
    """Credentials for whichever registry an image reference points at."""
    from services import update_service

    registry, _repo, _tag = update_service.parse_ref(image)
    return credentials_for(registry)


def configured_hosts() -> list[str]:
    with _lock:
        return sorted(_credentials)


# --------------------------------------------------------------------------- #
# The Docker CLI's config.json
# --------------------------------------------------------------------------- #


def docker_config_dir() -> str:
    return os.path.join(settings.DATA_DIR, "docker")


def _write_docker_config(creds: dict[str, tuple[str, str]]) -> None:
    """Write the auths file ``docker compose pull`` reads.

    The file is rewritten from the database every time, so removing a registry
    in the UI actually revokes the CLI's access rather than leaving a stale
    login behind.
    """
    directory = docker_config_dir()
    path = os.path.join(directory, "config.json")
    auths = {}
    for host, (username, password) in creds.items():
        key = DOCKER_HUB_CONFIG_KEY if host == DOCKER_HUB else host
        token = base64.b64encode(f"{username}:{password}".encode("utf-8")).decode("ascii")
        auths[key] = {"auth": token}
    try:
        os.makedirs(directory, mode=0o700, exist_ok=True)
        tmp = f"{path}.tmp"
        # Written 0600 before it is put in place, so the credentials are never
        # briefly world-readable.
        fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            json.dump({"auths": auths}, fh)
        os.replace(tmp, path)
    except OSError as exc:
        logger.warning("Could not write the Docker CLI credentials file: %s", exc)


def cli_env() -> dict:
    """Environment for a ``docker``/``docker compose`` subprocess.

    Points DOCKER_CONFIG at our generated file rather than writing into
    ``~/.docker``, so what StackPilot manages stays separate from anything the
    image ships or an operator put there by hand.
    """
    return {**os.environ, "DOCKER_CONFIG": docker_config_dir()}


# --------------------------------------------------------------------------- #
# Verifying credentials
# --------------------------------------------------------------------------- #


async def verify(host: str, username: str, password: str) -> None:
    """Check that a registry accepts these credentials. Raises RegistryError.

    Asks for a pull-scoped token the way a client would, and treats only an
    outright 401 as "wrong credentials" — a registry that answers anything else
    is reachable and talking, which is as much as a credentials check can
    honestly claim.
    """
    registry = canonical_host(host)
    url = f"https://{registry}/v2/"
    auth = (username, password)
    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=TIMEOUT) as client:
            response = await client.get(url, auth=auth)
            if response.status_code == 401:
                challenge = response.headers.get("WWW-Authenticate", "")
                if challenge.lower().startswith("bearer"):
                    token = await _token(client, challenge, auth)
                    if not token:
                        raise RegistryError("The registry rejected these credentials")
                    return
                raise RegistryError("The registry rejected these credentials")
            if response.status_code >= 500:
                raise RegistryError(
                    f"The registry answered {response.status_code}; try again later"
                )
    except httpx.HTTPError as exc:
        raise RegistryError(f"Could not reach {registry}: {exc}") from exc


async def _token(
    client: httpx.AsyncClient, challenge: str, auth: tuple[str, str]
) -> Optional[str]:
    """Follow a Bearer challenge with credentials attached."""
    params = {}
    for part in challenge[len("Bearer ") :].split(","):
        if "=" in part:
            key, value = part.split("=", 1)
            params[key.strip()] = value.strip().strip('"')
    realm = params.pop("realm", None)
    if not realm:
        return None
    try:
        response = await client.get(realm, params=params, auth=auth, timeout=TIMEOUT)
        if response.status_code == 401:
            return None
        response.raise_for_status()
        data = response.json()
        return data.get("token") or data.get("access_token")
    except (httpx.HTTPError, ValueError):
        return None