stackpilot/backend/services/compose_service.py

"""File-based stack storage and Docker Compose lifecycle.

The compose YAML on disk is always the source of truth. The database only
stores metadata (name, description, timestamps).
"""
from __future__ import annotations

import asyncio
import os
import re
import shutil
import stat
import tempfile
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

from config import settings
from docker_client import DockerError, get_client, safe_call
from services import registry_service

COMPOSE_FILENAMES = ("compose.yaml", "compose.yml", "docker-compose.yml", "docker-compose.yaml")
DEFAULT_COMPOSE_NAME = "compose.yaml"
COMPOSE_LABEL = "com.docker.compose.project"
SERVICE_LABEL = "com.docker.compose.service"


class StackFileError(Exception):
    pass


# --------------------------------------------------------------------------- #
# Slug / paths
# --------------------------------------------------------------------------- #


_PROJECT_INVALID_RE = re.compile(r"[^a-z0-9_-]+")


def project_name(stack_id: str) -> str:
    """The compose project name for a stack.

    The same normalisation compose applies to a directory name when it picks
    the project name itself (lowercase, only ``a-z0-9_-``, no leading ``-``/``_``).
    For every stack StackPilot creates that is the id unchanged. A folder found
    on disk as ``Nextcloud`` or ``my.app`` used to be passed as ``-p Nextcloud``,
    which compose refuses — so every action on such a stack failed, while
    ``docker compose up`` in that folder worked.
    """
    return _PROJECT_INVALID_RE.sub("", stack_id.lower()).lstrip("_-") or "stack"


def stack_ids_by_project(stack_ids) -> dict[str, str]:
    """``{project: stack_id}``, to key per-project Docker data by stack id."""
    return {project_name(sid): sid for sid in stack_ids}


def stack_ids_by_project_in(session) -> dict[str, str]:
    """:func:`stack_ids_by_project` for every stack in the database."""
    from sqlmodel import select

    from models.stack import Stack

    return stack_ids_by_project(s.id for s in session.exec(select(Stack)).all())


def owner_stack(project: Optional[str], ids: dict[str, str]) -> Optional[str]:
    """The stack id behind a ``com.docker.compose.project`` label, else the
    label itself (a project StackPilot does not manage)."""
    return ids.get(project, project) if project else project


def slugify(name: str) -> str:
    slug = re.sub(r"[^a-z0-9_-]+", "-", name.strip().lower())
    slug = re.sub(r"-{2,}", "-", slug).strip("-_")
    return slug or "stack"


def stacks_root(override: Optional[str] = None) -> str:
    return override or settings.STACKS_DIR


def stack_dir(stack_id: str, override: Optional[str] = None) -> str:
    return os.path.join(stacks_root(override), stack_id)


def find_compose_file(directory: str) -> Optional[str]:
    for name in COMPOSE_FILENAMES:
        candidate = os.path.join(directory, name)
        if os.path.isfile(candidate):
            return candidate
    return None


def compose_path(stack_id: str, override: Optional[str] = None) -> str:
    directory = stack_dir(stack_id, override)
    return find_compose_file(directory) or os.path.join(directory, DEFAULT_COMPOSE_NAME)


def env_path(stack_id: str, override: Optional[str] = None) -> str:
    return os.path.join(stack_dir(stack_id, override), ".env")


# --------------------------------------------------------------------------- #
# Read / write files
# --------------------------------------------------------------------------- #


def read_compose(stack_id: str, override: Optional[str] = None) -> str:
    path = compose_path(stack_id, override)
    if not os.path.isfile(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


def read_env(stack_id: str, override: Optional[str] = None) -> str:
    path = env_path(stack_id, override)
    if not os.path.isfile(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


#: Name parts that mark a value as a credential (DB_PASSWORD, JWT_SECRET,
#: CF_API_KEY, TS_AUTHKEY …), matched per ``_``/``.``/``-`` separated part so
#: that KEYCLOAK_ADMIN or a "monkey" service stay readable.
_SECRET_PARTS = {
    "pass", "passwd", "password", "passphrase", "pwd", "secret", "token", "key",
    "apikey", "authkey", "credential", "credentials", "private",
}
#: ``KEY: value`` and ``- KEY=value``, keeping indentation and quotes.
_MAPPING_LINE = re.compile(r"^(\s*)([\"']?)([A-Za-z0-9_.-]+)\2(\s*:\s+)(.+?)\s*$")
_LIST_LINE = re.compile(r"^(\s*-\s*)([\"']?)([A-Za-z0-9_.-]+)=(.*?)\2\s*$")
#: ``${VAR}``, ``$VAR`` or ``${VAR:?message}`` — but not ``${VAR:-default}``,
#: whose default is a value like any other.
_ONLY_A_VARIABLE = re.compile(
    r"^[\"']?\$(?:\{[A-Za-z_][A-Za-z0-9_]*(?::?\?[^}]*)?\}|[A-Za-z_][A-Za-z0-9_]*)[\"']?$"
)
MASK = "••••••"


def _secret_name(name: str) -> bool:
    return any(part in _SECRET_PARTS for part in re.split(r"[_.-]", name.lower()))


def mask_secrets(compose: str) -> str:
    """The compose file with credential values replaced by ``••••••``.

    For the read-only role, which is not given the ``.env`` — but plenty of
    compose files carry the password inline (``POSTGRES_PASSWORD: hunter2``),
    and the Compose tab showed it. A reference to a variable (``${DB_PASS}``)
    is not a secret and stays. Line-based on purpose: the file keeps its
    layout and comments, and a file that does not parse is masked all the same.
    """
    out = []
    for line in compose.splitlines(keepends=True):
        body = line.rstrip("\r\n")
        ending = line[len(body):]
        m = _LIST_LINE.match(body)
        if m and _secret_name(m.group(3)) and m.group(4) and not _ONLY_A_VARIABLE.match(m.group(4)):
            body = f"{m.group(1)}{m.group(2)}{m.group(3)}={MASK}{m.group(2)}"
        else:
            m = _MAPPING_LINE.match(body)
            if (
                m
                and _secret_name(m.group(3))
                and not m.group(5).startswith(("|", ">", "{", "[", "&", "*", "#"))
                and not _ONLY_A_VARIABLE.match(m.group(5))
            ):
                body = f"{m.group(1)}{m.group(2)}{m.group(3)}{m.group(2)}{m.group(4)}{MASK}"
        out.append(body + ending)
    return "".join(out)


def write_compose(stack_id: str, content: str, override: Optional[str] = None) -> None:
    directory = stack_dir(stack_id, override)
    os.makedirs(directory, exist_ok=True)
    path = compose_path(stack_id, override)
    # Non-destructive: back up existing file first.
    if os.path.isfile(path):
        shutil.copy2(path, path + ".bak")
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        fh.write(content)
    os.replace(tmp, path)


def write_env(stack_id: str, content: str, override: Optional[str] = None) -> None:
    directory = stack_dir(stack_id, override)
    os.makedirs(directory, exist_ok=True)
    path = env_path(stack_id, override)
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(content)


def _mount_points() -> set[str]:
    """Every mount point this process sees (from /proc/self/mountinfo)."""
    points: set[str] = set()
    try:
        with open("/proc/self/mountinfo", "r", encoding="utf-8") as fh:
            for line in fh:
                fields = line.split()
                if len(fields) >= 5:
                    # Spaces etc. are octal-escaped in mountinfo.
                    points.add(fields[4].encode().decode("unicode_escape"))
    except OSError:
        pass
    return points


def delete_stack_files(stack_id: str, override: Optional[str] = None) -> list[str]:
    """Remove the stack folder, but never what is mounted inside it.

    A stack folder can have a NAS share or a second disk mounted in it (a media
    library at ``./media``). ``rmtree`` would have walked straight into it and
    deleted the share's contents along with the compose file. Directories on
    another filesystem, or that are mount points of their own (a bind mount of
    the same disk), are left exactly as they are. Returns those paths.
    """
    directory = stack_dir(stack_id, override)
    if not os.path.isdir(directory):
        return []
    mounts = _mount_points()
    root_dev = os.lstat(directory).st_dev
    kept: list[str] = []

    def remove(path: str) -> None:
        with os.scandir(path) as it:
            entries = list(it)
        for entry in entries:
            if entry.is_dir(follow_symlinks=False):
                if (
                    entry.stat(follow_symlinks=False).st_dev != root_dev
                    or entry.path in mounts
                ):
                    kept.append(entry.path)
                    continue
                remove(entry.path)
            else:
                os.unlink(entry.path)
        try:
            os.rmdir(path)
        except OSError:
            pass  # still holds something that was kept

    remove(directory)
    return kept


def _skip_special_files(directory: str, names: list[str]) -> list[str]:
    """copytree ignore hook: sockets, FIFOs and devices cannot be copied."""
    skipped = []
    for name in names:
        try:
            mode = os.lstat(os.path.join(directory, name)).st_mode
        except OSError:
            continue
        if not (stat.S_ISDIR(mode) or stat.S_ISREG(mode) or stat.S_ISLNK(mode)):
            skipped.append(name)
    return skipped


def clone_stack_files(src_id: str, dst_id: str, override: Optional[str] = None) -> None:
    src = stack_dir(src_id, override)
    dst = stack_dir(dst_id, override)
    if os.path.isdir(dst):
        raise StackFileError(f"Target stack '{dst_id}' already exists")
    try:
        # symlinks=True copies links as links: following them would duplicate
        # whatever they point at (or fail on a dangling one). A container's
        # runtime socket in a data folder would fail the whole copy.
        shutil.copytree(src, dst, symlinks=True, ignore=_skip_special_files)
    except (OSError, shutil.Error) as exc:
        shutil.rmtree(dst, ignore_errors=True)
        raise StackFileError(f"Could not copy the stack folder: {exc}") from exc


def force_remove_project(stack_id: str) -> int:
    """Remove a project's containers and networks through the Docker API.

    The fallback for deleting a stack whose ``compose down`` failed — typically
    because the compose file no longer parses. Without it the containers kept
    running with nothing left in StackPilot to stop them. Returns how many
    containers went.
    """
    client = get_client()
    label = f"{COMPOSE_LABEL}={project_name(stack_id)}"
    removed = 0
    for container in safe_call(client.containers.list, all=True, filters={"label": label}):
        try:
            container.remove(force=True)
            removed += 1
        except Exception:  # noqa: BLE001 - best effort, report what went
            pass
    try:
        for network in safe_call(client.networks.list, filters={"label": label}):
            try:
                network.remove()
            except Exception:  # noqa: BLE001
                pass
    except DockerError:
        pass
    return removed


def discover_stacks(override: Optional[str] = None) -> list[str]:
    """Return ids of all directories under STACKS_DIR that contain a compose file."""
    root = stacks_root(override)
    if not os.path.isdir(root):
        return []
    found = []
    for entry in sorted(os.listdir(root)):
        # Hidden directories are StackPilot's own scratch space (a restore
        # stages there before swapping it in), never a stack.
        if entry.startswith("."):
            continue
        directory = os.path.join(root, entry)
        if os.path.isdir(directory) and find_compose_file(directory):
            found.append(entry)
    return found


# --------------------------------------------------------------------------- #
# Live status from Docker
# --------------------------------------------------------------------------- #


@dataclass
class ContainerInfo:
    id: str
    name: str
    service: str
    image: str
    state: str  # running, exited, ...
    status: str  # human string
    health: Optional[str] = None
    ports: list[dict] = field(default_factory=list)
    created: Optional[str] = None


def _parse_ports(attrs: dict) -> list[dict]:
    ports = []
    bindings = (attrs.get("NetworkSettings") or {}).get("Ports") or {}
    for container_port, host in (bindings or {}).items():
        if host:
            for binding in host:
                ports.append(
                    {
                        "container": container_port,
                        "host_ip": binding.get("HostIp"),
                        "host_port": binding.get("HostPort"),
                    }
                )
        else:
            ports.append({"container": container_port, "host_port": None})
    return ports


def containers_for_stack(stack_id: str) -> list[ContainerInfo]:
    client = get_client()
    raw = safe_call(
        client.containers.list,
        all=True,
        filters={"label": f"{COMPOSE_LABEL}={project_name(stack_id)}"},
    )
    result = []
    for c in raw:
        attrs = c.attrs
        state = attrs.get("State", {}) or {}
        health = (state.get("Health") or {}).get("Status")
        result.append(
            ContainerInfo(
                id=c.id,
                name=c.name,
                service=c.labels.get(SERVICE_LABEL, c.name),
                # Use the configured image name from the inspect we already have;
                # reading c.image.tags would trigger a separate image-inspect call
                # per container.
                image=attrs.get("Config", {}).get("Image", "") or attrs.get("Image", ""),
                state=state.get("Status", c.status),
                status=attrs.get("State", {}).get("Status", c.status),
                health=health,
                ports=_parse_ports(attrs),
                created=attrs.get("Created"),
            )
        )
    return result


def _status_from_states(states: list[str]) -> str:
    if not states:
        return "stopped"
    if any(s == "dead" for s in states):
        return "error"
    running = [s for s in states if s == "running"]
    if len(running) == len(states):
        return "running"
    if running:
        return "partial"
    return "stopped"


def compute_status(stack_id: str, containers: Optional[list[ContainerInfo]] = None) -> str:
    """Status for one stack, from its containers alone.

    "updating" is not derived here: whether an operation is in flight lives in
    ``stack_lock_service``, and callers that want to show it overlay the lock on
    top of this. Pass already-fetched ``containers`` to avoid a redundant Docker
    round-trip (the detail view already has them).
    """
    try:
        if containers is None:
            containers = containers_for_stack(stack_id)
    except DockerError:
        return "unknown"
    return _status_from_states([c.state for c in containers])


def stack_status_summaries() -> dict[str, dict]:
    """One cheap Docker call → ``{project: {status, total, running}}`` for every
    stack at once.

    Uses the low-level container *summary* list (``GET /containers/json``, no
    per-container inspect) grouped by the compose project label. This is far
    cheaper than calling :func:`containers_for_stack` (which full-inspects every
    container) once per stack for the overview — turning the stacks list from
    O(stacks × containers) Docker round-trips into a single one.
    """
    client = get_client()
    raw = safe_call(client.api.containers, all=True)
    by_project: dict[str, list[str]] = {}
    for c in raw:
        project = (c.get("Labels") or {}).get(COMPOSE_LABEL)
        if project:
            by_project.setdefault(project, []).append(c.get("State", ""))
    return {
        project: {
            "status": _status_from_states(states),
            "total": len(states),
            "running": sum(1 for s in states if s == "running"),
        }
        for project, states in by_project.items()
    }


# --------------------------------------------------------------------------- #
# Compose CLI lifecycle (async subprocess)
# --------------------------------------------------------------------------- #


def _compose_base_cmd(stack_id: str, override: Optional[str] = None) -> list[str]:
    directory = stack_dir(stack_id, override)
    compose_file = find_compose_file(directory) or os.path.join(directory, DEFAULT_COMPOSE_NAME)
    return [
        "docker",
        "compose",
        "-p",
        project_name(stack_id),
        "--project-directory",
        directory,
        "-f",
        compose_file,
    ]


async def run_compose(
    stack_id: str,
    args: list[str],
    override: Optional[str] = None,
    timeout: float = 600.0,
) -> dict:
    """Run a `docker compose` subcommand. Returns {returncode, stdout, stderr}."""
    cmd = _compose_base_cmd(stack_id, override) + args
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        env=registry_service.cli_env(),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError as exc:
        proc.kill()
        await proc.wait()
        raise StackFileError(f"compose command timed out: {' '.join(args)}") from exc
    return {
        "returncode": proc.returncode,
        "stdout": stdout_b.decode("utf-8", "replace"),
        "stderr": stderr_b.decode("utf-8", "replace"),
        "command": " ".join(args),
    }


async def validate_yaml(
    content: str, env_content: str = "", project_dir: Optional[str] = None
) -> dict:
    """Validate compose YAML by running ``docker compose config -q`` against it.
    Returns ``{"ok": bool, "errors": str}`` — the errors string carries
    compose's own diagnostics (parse errors, unknown keys, bad interpolation)
    so the editor can surface them before saving.

    With ``project_dir`` (an existing stack's folder) relative paths resolve
    against that folder, as they will on deploy. Validated from a bare temp
    directory, a stack using ``env_file: ./app.env`` — or an ``include``, a
    ``secrets: file:`` — was reported broken although it deploys fine. The
    edited ``.env`` is still the one used for interpolation.
    """
    with tempfile.TemporaryDirectory(prefix="sp-validate-") as tmp:
        compose_file = os.path.join(tmp, DEFAULT_COMPOSE_NAME)
        with open(compose_file, "w", encoding="utf-8") as fh:
            fh.write(content)
        env_file = os.path.join(tmp, ".env")
        with open(env_file, "w", encoding="utf-8") as fh:
            fh.write(env_content or "")
        directory = project_dir if project_dir and os.path.isdir(project_dir) else tmp
        cmd = [
            "docker", "compose",
            "--project-directory", directory,
            "-f", compose_file,
            "--env-file", env_file,
            "config", "-q",
        ]
        if directory != tmp:
            # The folder's name may not be a valid project name; config does
            # not care what the project is called.
            cmd[2:2] = ["-p", "sp-validate"]
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            env=registry_service.cli_env(),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            _out, err = await asyncio.wait_for(proc.communicate(), timeout=30.0)
        except asyncio.TimeoutError as exc:
            proc.kill()
            await proc.wait()
            raise StackFileError("compose config validation timed out") from exc
        ok = proc.returncode == 0
        errors = "" if ok else err.decode("utf-8", "replace").strip()
        # Name the file the user is editing, not our scratch copy of it.
        errors = errors.replace(compose_file, DEFAULT_COMPOSE_NAME).replace(env_file, ".env")
        errors = errors.replace(tmp + os.sep, "./")
        return {"ok": ok, "errors": errors}


async def stream_compose(
    stack_id: str, args: list[str], override: Optional[str] = None
):
    """Yield lines from a `docker compose` subcommand as they are produced.

    The subprocess is killed when the consumer stops early (the generator is
    closed or its task cancelled). For ``logs -f`` that is the only way it ever
    ends: without this every log view that was closed left a ``compose logs``
    process behind for as long as the containers lived.
    """
    cmd = _compose_base_cmd(stack_id, override) + args
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        env=registry_service.cli_env(),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    assert proc.stdout is not None
    try:
        async for raw in proc.stdout:
            yield raw.decode("utf-8", "replace").rstrip("\n")
        await proc.wait()
    finally:
        if proc.returncode is None:
            try:
                proc.kill()
            except ProcessLookupError:
                pass
            # Reap it, so it does not linger as a zombie. Shielded: this runs
            # while the consuming task is being cancelled.
            try:
                await asyncio.shield(proc.wait())
            except (asyncio.CancelledError, ProcessLookupError):
                pass


_json_progress: Optional[bool] = None


async def supports_json_progress() -> bool:
    """Whether this Docker Compose understands ``--progress json``.

    The JSON progress stream carries per-layer ``current``/``total`` bytes, which
    the deploy console turns into a real progress bar. Older compose releases
    reject the value, so probe once (cheap, no side effects) and cache it; on a
    negative result callers fall back to the plain text stream.
    """
    global _json_progress
    if _json_progress is None:
        try:
            proc = await asyncio.create_subprocess_exec(
                "docker", "compose", "--progress", "json", "version",
                env=registry_service.cli_env(),
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
            )
            rc = await asyncio.wait_for(proc.wait(), timeout=15.0)
            _json_progress = rc == 0
        except Exception:  # noqa: BLE001 - probe failure just disables the feature
            _json_progress = False
    return _json_progress


async def _stream_phase(
    stack_id: str, args: list[str], override: Optional[str], json_progress: bool
):
    """One compose subcommand, streamed. Yields ``("log", line)`` per output
    line, then ``("rc", returncode)`` exactly once."""
    cmd = _compose_base_cmd(stack_id, override)
    if json_progress:
        # Global flag, must precede the subcommand.
        cmd += ["--progress", "json"]
    cmd += args
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        env=registry_service.cli_env(),
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    assert proc.stdout is not None
    async for raw in proc.stdout:
        yield ("log", raw.decode("utf-8", "replace").rstrip("\n"))
    await proc.wait()
    yield ("rc", proc.returncode)


async def stream_up(stack_id: str, override: Optional[str] = None):
    """Run `compose up -d` streaming combined output, so the deploy console can
    show image-pull and container-create progress live.

    Yields ``("log", line)`` for each output line, then ``("done", returncode)``.
    """
    json_progress = await supports_json_progress()
    async for kind, payload in _stream_phase(
        stack_id, ["up", "-d", "--remove-orphans"], override, json_progress
    ):
        yield ("done", payload) if kind == "rc" else ("log", payload)


async def stream_update(stack_id: str, override: Optional[str] = None):
    """Run `compose pull` then `compose up -d`, streaming both phases, so the
    stacks list can show real update progress instead of a spinner.

    Yields ``("log", line)`` for each output line of either phase, then
    ``("done", returncode)`` once. A failed pull short-circuits: recreating
    containers on images that never came down would only make things worse.
    """
    json_progress = await supports_json_progress()
    rc = 0
    for args in (["pull"], ["up", "-d", "--remove-orphans"]):
        async for kind, payload in _stream_phase(stack_id, args, override, json_progress):
            if kind == "log":
                yield ("log", payload)
            else:
                rc = payload
        if rc != 0:
            break
    yield ("done", rc)


# Convenience lifecycle wrappers ------------------------------------------------


async def up(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["up", "-d", "--remove-orphans"], override)


async def down(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["down"], override)


async def start(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["start"], override)


async def stop(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["stop"], override)


async def restart(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["restart"], override)


async def pull(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["pull"], override)


async def update(stack_id: str, override: Optional[str] = None) -> dict:
    """Pull then up -d.

    A failed pull stops here, like the streamed update does. Running ``up``
    anyway recreated nothing new — the old images were still there — and its
    exit code 0 then reported the update as a success.
    """
    pull_res = await run_compose(stack_id, ["pull"], override)
    if pull_res["returncode"] != 0:
        return {**pull_res, "command": "pull + up -d"}
    up_res = await run_compose(stack_id, ["up", "-d", "--remove-orphans"], override)
    return {
        "returncode": up_res["returncode"],
        "stdout": pull_res["stdout"] + "\n" + up_res["stdout"],
        "stderr": pull_res["stderr"] + "\n" + up_res["stderr"],
        "command": "pull + up -d",
    }


async def logs(
    stack_id: str,
    service: Optional[str] = None,
    tail: int = 200,
    override: Optional[str] = None,
) -> dict:
    args = ["logs", "--no-color", "--tail", str(tail), "--timestamps"]
    if service:
        # "--" so a service name can never be read as a compose option.
        args += ["--", service]
    return await run_compose(stack_id, args, override, timeout=60.0)


def now() -> datetime:
    return datetime.now(timezone.utc)