Unblu MCP Server

import asyncio
import shutil
import socket
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from fastmcp.utilities.logging import get_logger

from unblu_mcp._internal.exceptions import ConfigurationError
from unblu_mcp._internal.providers import ConnectionConfig, ConnectionProvider

_logger = get_logger(__name__)

if TYPE_CHECKING:
    from typing import Any


@dataclass
class K8sEnvironmentConfig:
    """Configuration for a Kubernetes environment."""

    name: str
    """Environment name (e.g., dev, staging, prod)."""

    local_port: int
    """Local port for port-forwarding."""

    namespace: str
    """Kubernetes namespace."""

    service: str = "haproxy"
    """Kubernetes service name."""

    service_port: int = 8080
    """Service port to forward."""

    api_path: str = "/app/rest/v4"
    """API path prefix."""


# Config file paths (relative to project root)
_CONFIG_DIR = Path(__file__).parent.parent.parent.parent / "config"
_USER_CONFIG = _CONFIG_DIR / "k8s_environments.yaml"
_EXAMPLE_CONFIG = _CONFIG_DIR / "k8s_environments.example.yaml"


def _load_environments_from_yaml(path: Path) -> dict[str, K8sEnvironmentConfig]:
    """Load environment configurations from a YAML file."""
    try:
        import yaml  # noqa: PLC0415
    except ImportError as e:
        msg = "PyYAML is required for K8s environments. Install with: pip install pyyaml"
        raise ImportError(msg) from e

    if not path.exists():
        return {}

    with path.open(encoding="utf-8") as f:
        data: dict[str, Any] = yaml.safe_load(f) or {}

    environments: dict[str, K8sEnvironmentConfig] = {}
    for name, config in data.get("environments", {}).items():
        environments[name] = K8sEnvironmentConfig(
            name=name,
            local_port=config["local_port"],
            namespace=config["namespace"],
            service=config.get("service", "haproxy"),
            service_port=config.get("service_port", 8080),
            api_path=config.get("api_path", "/app/rest/v4"),
        )
    return environments


def _get_default_environments() -> dict[str, K8sEnvironmentConfig]:
    """Get environment configurations from user config or example file."""
    # Try user config first (gitignored, may contain sensitive data)
    if _USER_CONFIG.exists():
        return _load_environments_from_yaml(_USER_CONFIG)
    # Fall back to example config
    return _load_environments_from_yaml(_EXAMPLE_CONFIG)


class K8sConnectionProvider(ConnectionProvider):
    """Connection provider for Kubernetes deployments using port-forwarding.

    This provider:
    - Automatically starts kubectl port-forward on setup
    - Cleans up the port-forward process on teardown
    - Configures trusted headers authentication
    - Supports multiple environments with different ports

    Args:
        environment: Environment name (dev, staging, prod) or custom K8sEnvironmentConfig.
        trusted_user_id: User ID for trusted headers auth (default: superadmin).
        trusted_user_role: User role for trusted headers auth (default: SUPER_ADMIN).
        environments: Custom environment configurations (overrides defaults).

    Example:
        provider = K8sConnectionProvider(environment="dev")
        await provider.setup()  # Starts port-forward
        config = provider.get_config()  # Returns connection config
        await provider.teardown()  # Stops port-forward
    """

    def __init__(
        self,
        environment: str | K8sEnvironmentConfig = "dev",
        trusted_user_id: str = "superadmin",
        trusted_user_role: str = "SUPER_ADMIN",
        environments: dict[str, K8sEnvironmentConfig] | None = None,
    ) -> None:
        self._environments = environments or _get_default_environments()

        if isinstance(environment, str):
            if environment not in self._environments:
                valid = ", ".join(self._environments.keys())
                msg = f"Unknown environment '{environment}'. Valid: {valid}"
                raise ValueError(msg)
            self._env_config = self._environments[environment]
        else:
            self._env_config = environment

        self._trusted_user_id = trusted_user_id
        self._trusted_user_role = trusted_user_role
        self._port_forward_process: subprocess.Popen[bytes] | None = None
        self._owns_port_forward = False  # Whether we started the port-forward

    @property
    def environment(self) -> str:
        """Return the current environment name."""
        return self._env_config.name

    @property
    def local_port(self) -> int:
        """Return the local port for this environment."""
        return self._env_config.local_port

    async def setup(self) -> None:
        """Start kubectl port-forward if not already running.

        Uses a simple port-check approach: if the port is already in use,
        assume another instance is handling it. Otherwise, start our own.
        """
        if self._is_port_in_use():
            # Port already in use - another instance or external process has it
            _logger.debug(
                "Port %d already in use, reusing existing connection",
                self._env_config.local_port,
            )
            self._owns_port_forward = False
        else:
            # Port not in use - we need to start port-forward
            _logger.debug("Port %d not in use, starting port-forward", self._env_config.local_port)
            self._owns_port_forward = True
            await self._start_port_forward()

    async def _start_port_forward(self) -> None:
        """Start the kubectl port-forward process."""
        if not shutil.which("kubectl"):
            msg = "kubectl not found in PATH. Install kubectl to use the K8s provider."
            raise ConfigurationError(msg)

        # Check if kubectl is authenticated (with timeout to avoid hanging)
        try:
            auth_check = subprocess.run(  # noqa: S603
                ["kubectl", "auth", "can-i", "get", "pods", "-n", self._env_config.namespace],  # noqa: S607
                capture_output=True,
                text=True,
                check=False,
                timeout=5,  # 5 second timeout to avoid hanging on OIDC prompts
            )
        except subprocess.TimeoutExpired:
            msg = (
                f"kubectl auth check timed out for namespace '{self._env_config.namespace}'. "
                f"This usually means kubectl is waiting for authentication (e.g., OIDC login). "
                f"Please authenticate using your cluster's auth method (e.g., cloud CLI, kubelogin, or kubeconfig)"
            )
            raise ConfigurationError(msg) from None
        if auth_check.returncode != 0:
            stderr = auth_check.stderr.strip()
            msg = (
                f"kubectl is not authenticated or lacks permissions for namespace '{self._env_config.namespace}'. "
                f"Please authenticate using your cluster's auth method (e.g., cloud CLI, kubelogin). "
                f"Error: {stderr or 'unknown'}"
            )
            raise ConfigurationError(msg)

        cmd = [
            "kubectl",
            "port-forward",
            "-n",
            self._env_config.namespace,
            f"svc/{self._env_config.service}",
            f"{self._env_config.local_port}:{self._env_config.service_port}",
        ]

        _logger.debug("Starting port-forward: %s", " ".join(cmd))

        self._port_forward_process = subprocess.Popen(  # noqa: S603
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        # Wait for port to become available
        await self._wait_for_port()

    async def _wait_for_port(self, timeout: float = 10.0) -> None:
        """Wait for the port to become available after starting port-forward."""
        iterations = int(timeout / 0.5)
        for _ in range(iterations):
            await asyncio.sleep(0.5)
            if self._is_port_in_use():
                _logger.debug("Port %d is now available", self._env_config.local_port)
                return

            # Check if process died early
            if self._port_forward_process is not None:
                retcode = self._port_forward_process.poll()
                if retcode is not None:
                    # Process exited - get the error
                    _, stderr = self._port_forward_process.communicate()
                    stderr_text = stderr.decode().strip() if stderr else "unknown error"
                    self._port_forward_process = None
                    msg = (
                        f"kubectl port-forward failed for {self._env_config.name}: {stderr_text}. "
                        f"Ensure you are authenticated to the K8s cluster and have access to namespace '{self._env_config.namespace}'."
                    )
                    raise ConfigurationError(msg)

        # Timeout - clean up the process we started
        if self._port_forward_process is not None:
            self._port_forward_process.kill()
            _, stderr = self._port_forward_process.communicate()
            stderr_text = stderr.decode().strip() if stderr else ""
            self._port_forward_process = None
            msg = (
                f"Port-forward timed out for {self._env_config.name} - port did not become available. "
                f"kubectl stderr: {stderr_text or 'none'}. "
                f"Ensure kubectl is authenticated and the service '{self._env_config.service}' exists in namespace '{self._env_config.namespace}'."
            )
            raise ConfigurationError(msg)

    async def teardown(self) -> None:
        """Stop the port-forward process only if we started it."""
        if self._owns_port_forward and self._port_forward_process is not None:
            _logger.debug("Stopping port-forward for %s", self._env_config.name)
            self._port_forward_process.terminate()
            try:
                self._port_forward_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._port_forward_process.kill()
            self._port_forward_process = None
        elif not self._owns_port_forward:
            _logger.debug("Not owner, skipping port-forward cleanup for %s", self._env_config.name)

    async def ensure_connection(self) -> None:
        """Ensure the port-forward is running, restarting if needed.

        Call this before making API requests to handle cases where:
        - Our port-forward process died
        - An external port-forward we were using is no longer available
        """
        if self._is_port_in_use():
            return  # Port is available, nothing to do

        # Port is not available - need to (re)start port-forward
        if self._owns_port_forward and self._port_forward_process is not None:
            # Clean up existing process (dead or malfunctioning)
            retcode = self._port_forward_process.poll()
            if retcode is not None:
                _logger.warning(
                    "Port-forward process died (exit code %d), restarting...",
                    retcode,
                )
            else:
                _logger.warning("Port-forward process alive but port not available, killing and restarting...")
                self._port_forward_process.kill()
                try:
                    self._port_forward_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    _logger.warning("Process did not terminate after kill, continuing anyway")
            self._port_forward_process = None

        # Start a new port-forward
        _logger.info("Restarting port-forward for %s", self._env_config.name)
        self._owns_port_forward = True
        await self._start_port_forward()

    def get_config(self) -> ConnectionConfig:
        """Return connection config with trusted headers."""
        return ConnectionConfig(
            base_url=f"http://localhost:{self._env_config.local_port}{self._env_config.api_path}",
            headers={
                "x-unblu-trusted-user-id": self._trusted_user_id,
                "x-unblu-trusted-user-role": self._trusted_user_role,
            },
        )

    async def health_check(self) -> bool:
        """Check if the port-forward is healthy."""
        return self._is_port_in_use()

    def _is_port_in_use(self) -> bool:
        """Check if the local port is in use."""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            return s.connect_ex(("localhost", self._env_config.local_port)) == 0


def detect_environment_from_context() -> str | None:
    """Detect the environment from the current kubectl context.

    Matches environment names from the loaded configuration against
    patterns in the current kubectl context name.

    Returns:
        Environment name or None if not detected.
    """
    try:
        result = subprocess.run(
            ["kubectl", "config", "current-context"],  # noqa: S607
            capture_output=True,
            text=True,
            check=True,
            timeout=5,  # 5 second timeout to avoid hanging
        )
        context = result.stdout.strip()

        # Match against configured environment names
        environments = _get_default_environments()
        for env in environments:
            if f"-{env}-" in context or context.endswith(f"-{env}"):
                return env
        return None  # noqa: TRY300
    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
        return None