stackpilot/backend/services/icon_service.py

"""Stack icons — validation of the stored choice and custom-image storage.

A stack's ``icon`` column holds one of three things:

``None``
    Automatic. Nothing is stored and the UI derives an icon from the stack's
    name (see ``frontend/src/lib/stackIcons.ts``). Every stack that existed
    before this feature lands here, so an upgraded install shows sensible icons
    without a data migration.

``lucide:<name>``
    A built-in icon the user picked explicitly. The catalog of names lives in
    the frontend because that is the only place that can actually *render* one;
    keeping a second copy here would only add a list to drift out of sync. An
    unknown name is therefore not an error — the UI falls back to the automatic
    icon for it.

``custom:<ext>:<version>``
    An uploaded image at ``${DATA_DIR}/stack-icons/<stack_id>.<ext>``.
    ``<version>`` is the upload's unix timestamp. It carries no meaning beyond
    changing the column value on every re-upload, which is what makes the
    frontend's cache key change and the new image appear instead of the one the
    browser already has.
"""
from __future__ import annotations

import logging
import os
import re
import shutil
import time
from typing import Optional

from config import settings

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

#: Uploads are meant to be small app logos. The cap is deliberately generous
#: for a logo and still far too small to make the data directory grow.
MAX_ICON_BYTES = 512 * 1024

#: Extension per accepted image type. The extension is derived from the bytes
#: (see :func:`_sniff`), never from the upload's filename or Content-Type — a
#: client is free to lie about both.
_EXTENSIONS = {"png", "jpg", "gif", "webp", "svg"}

_BUILTIN_RE = re.compile(r"^lucide:[a-z0-9-]{1,48}$")
_CUSTOM_RE = re.compile(r"^custom:(png|jpg|gif|webp|svg):(\d{1,12})$")
_SAFE_ID_RE = re.compile(r"^[A-Za-z0-9._-]{1,128}$")

#: ``<img src>`` never executes script, but a custom icon is also reachable
#: directly under /api/..., where an SVG *would* run in the browser's own
#: context. Serving it as a download-only attachment keeps that door shut.
_CONTENT_TYPES = {
    "png": "image/png",
    "jpg": "image/jpeg",
    "gif": "image/gif",
    "webp": "image/webp",
    "svg": "image/svg+xml",
}


class IconError(Exception):
    """An icon value or upload the server refuses."""


# --------------------------------------------------------------------------- #
# Paths
# --------------------------------------------------------------------------- #


def icon_dir() -> str:
    return os.path.join(settings.DATA_DIR, "stack-icons")


def _safe_id(stack_id: str) -> str:
    """Reject anything that could escape the icon directory.

    Stack ids are slugs, so this never fires in practice — it is here because
    the id arrives from the URL and is about to be pasted into a filesystem
    path.
    """
    if not _SAFE_ID_RE.match(stack_id) or stack_id in (".", ".."):
        raise IconError(f"Invalid stack id '{stack_id}'")
    return stack_id


def custom_path(stack_id: str, ext: str) -> str:
    return os.path.join(icon_dir(), f"{_safe_id(stack_id)}.{ext}")


def custom_ext(value: Optional[str]) -> Optional[str]:
    """The file extension of a ``custom:`` icon value, or None for the rest."""
    match = _CUSTOM_RE.match(value or "")
    return match.group(1) if match else None


def content_type(ext: str) -> str:
    return _CONTENT_TYPES.get(ext, "application/octet-stream")


def file_for(stack_id: str, value: Optional[str]) -> Optional[str]:
    """Existing file backing a ``custom:`` icon value, or None."""
    ext = custom_ext(value)
    if not ext:
        return None
    path = custom_path(stack_id, ext)
    return path if os.path.isfile(path) else None


# --------------------------------------------------------------------------- #
# Stored value
# --------------------------------------------------------------------------- #


def normalize_choice(value: str) -> Optional[str]:
    """Validate an icon chosen through the API.

    An empty string means "back to automatic" and maps to ``None``. Only the
    built-in form is accepted here: a ``custom:`` value is minted by
    :func:`store_upload` and never taken from a client, or a stack could be
    pointed at another stack's icon file.
    """
    value = (value or "").strip()
    if not value:
        return None
    if _BUILTIN_RE.match(value):
        return value
    raise IconError(
        "Icon must be empty (automatic) or 'lucide:<name>'; upload custom "
        "images through POST /api/stacks/{id}/icon"
    )


# --------------------------------------------------------------------------- #
# Uploads
# --------------------------------------------------------------------------- #


def _sniff(data: bytes) -> str:
    """Extension for the image these bytes actually are.

    Trusting the declared Content-Type would mean storing (and later serving)
    whatever a client cares to send under an image's name.
    """
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        return "png"
    if data.startswith(b"\xff\xd8\xff"):
        return "jpg"
    if data.startswith(b"GIF87a") or data.startswith(b"GIF89a"):
        return "gif"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "webp"
    head = data[:512].lstrip()
    if head.startswith(b"<?xml") or head.startswith(b"<svg") or b"<svg" in head:
        return "svg"
    raise IconError("Unsupported image type — use PNG, JPEG, GIF, WebP or SVG")


def store_upload(stack_id: str, data: bytes) -> str:
    """Write an uploaded icon and return the value for ``Stack.icon``."""
    if not data:
        raise IconError("The uploaded file is empty")
    if len(data) > MAX_ICON_BYTES:
        raise IconError(
            f"Icon is too large ({len(data) // 1024} KiB); "
            f"the limit is {MAX_ICON_BYTES // 1024} KiB"
        )
    ext = _sniff(data)
    os.makedirs(icon_dir(), exist_ok=True)
    # A re-upload in a different format would otherwise leave the old file
    # behind as an orphan nothing ever cleans up.
    remove(stack_id)
    path = custom_path(stack_id, ext)
    with open(path, "wb") as fh:
        fh.write(data)
    return f"custom:{ext}:{int(time.time())}"


def remove(stack_id: str) -> None:
    """Delete every custom icon file belonging to a stack. Best effort."""
    for ext in _EXTENSIONS:
        try:
            os.remove(custom_path(stack_id, ext))
        except FileNotFoundError:
            continue
        except OSError as exc:  # noqa: PERF203 - one bad file must not block the rest
            logger.warning("Could not remove icon %s.%s: %s", stack_id, ext, exc)


def copy(src_id: str, dst_id: str, value: Optional[str]) -> Optional[str]:
    """Copy a stack's custom icon to another stack (used when cloning).

    Returns the icon value for the new stack: the copied ``custom:`` value, the
    unchanged built-in choice, or None when there is nothing to carry over.
    """
    ext = custom_ext(value)
    if not ext:
        return value
    source = file_for(src_id, value)
    if not source:
        return None
    os.makedirs(icon_dir(), exist_ok=True)
    try:
        shutil.copyfile(source, custom_path(dst_id, ext))
    except OSError as exc:
        logger.warning("Could not copy icon %s -> %s: %s", src_id, dst_id, exc)
        return None
    return f"custom:{ext}:{int(time.time())}"