haproxy-mcp/haproxy_mcp/tools/configuration.py

"""Configuration management tools for HAProxy MCP Server."""

import time

from ..config import (
    STATE_FILE,
    HAPROXY_CONTAINER,
    SUBPROCESS_TIMEOUT,
    STARTUP_RETRY_COUNT,
    logger,
)
from ..exceptions import HaproxyError
from ..haproxy_client import haproxy_cmd, haproxy_cmd_batch, reload_haproxy
from ..file_ops import (
    atomic_write_file,
    load_servers_config,
    get_domain_backend,
    get_backend_and_prefix,
)
from ..ssh_ops import run_command


def restore_servers_from_config() -> int:
    """Restore all servers from configuration file.

    Uses batched commands for efficiency - single TCP connection for all servers.

    Returns:
        Number of servers restored
    """
    config = load_servers_config()
    if not config:
        return 0

    # Build batch of all commands
    commands: list[str] = []
    server_info_list: list[tuple[str, str]] = []  # For logging on failure

    for domain, slots in config.items():
        backend = get_domain_backend(domain)
        if not backend:
            continue

        try:
            _, server_prefix = get_backend_and_prefix(domain)
        except ValueError as e:
            logger.warning("Invalid domain '%s': %s", domain, e)
            continue

        for slot_str, server_info in slots.items():
            try:
                slot = int(slot_str)
            except ValueError:
                logger.warning("Invalid slot '%s' for %s, skipping", slot_str, domain)
                continue

            ip = server_info.get("ip", "")
            if not ip:
                continue

            try:
                port = int(server_info.get("http_port", 80))
            except (ValueError, TypeError):
                logger.warning("Invalid port for %s slot %d, skipping", domain, slot)
                continue

            server = f"{server_prefix}_{slot}"
            commands.append(f"set server {backend}/{server} addr {ip} port {port}")
            commands.append(f"set server {backend}/{server} state ready")
            server_info_list.append((backend, server))

    if not commands:
        return 0

    # Execute all commands
    try:
        haproxy_cmd_batch(commands)
        return len(server_info_list)
    except HaproxyError as e:
        logger.warning("Batch restore failed: %s", e)
        # Fallback: try individual commands
        restored = 0
        for i in range(0, len(commands), 2):
            try:
                haproxy_cmd_batch([commands[i], commands[i + 1]])
                restored += 1
            except HaproxyError as e2:
                backend, server = server_info_list[i // 2]
                logger.warning("Failed to restore %s/%s: %s", backend, server, e2)
        return restored


def startup_restore() -> None:
    """Restore servers and certificates from config files on startup."""
    # Wait for HAProxy to be ready
    for _ in range(STARTUP_RETRY_COUNT):
        try:
            haproxy_cmd("show info")
            break
        except HaproxyError:
            time.sleep(1)
    else:
        logger.warning("HAProxy not ready, skipping restore")
        return

    # Restore servers
    try:
        count = restore_servers_from_config()
        if count > 0:
            logger.info("Restored %d servers from config", count)
    except (HaproxyError, OSError, ValueError) as e:
        logger.warning("Failed to restore servers: %s", e)

    # Restore certificates
    try:
        from .certificates import restore_certificates
        cert_count = restore_certificates()
        if cert_count > 0:
            logger.info("Restored %d certificates from config", cert_count)
    except (HaproxyError, IOError, OSError, ValueError) as e:
        logger.warning("Failed to restore certificates: %s", e)


def register_config_tools(mcp):
    """Register configuration management tools with MCP server."""

    @mcp.tool()
    def haproxy_reload() -> str:
        """Reload HAProxy configuration (validates config first).

        After reload, automatically restores server configurations from servers.json.

        Returns:
            Success message with restored server count, or error details if failed
        """
        success, msg = reload_haproxy()
        if not success:
            return msg

        # Wait for HAProxy to fully reload (new process takes over)
        # USR2 signal spawns new process but old one may still be serving
        time.sleep(2)

        # Verify HAProxy is responding
        for _ in range(STARTUP_RETRY_COUNT):
            try:
                haproxy_cmd("show info")
                break
            except HaproxyError:
                time.sleep(0.5)
        else:
            return "HAProxy reloaded but not responding after reload"

        # Restore servers from config after reload
        try:
            restored = restore_servers_from_config()
            return f"HAProxy configuration reloaded successfully ({restored} servers restored)"
        except (HaproxyError, IOError, OSError, ValueError) as e:
            logger.error("Failed to restore servers after reload: %s", e)
            return f"HAProxy reloaded but server restore failed: {e}"

    @mcp.tool()
    def haproxy_check_config() -> str:
        """Validate HAProxy configuration file syntax.

        Returns:
            Validation result or error details
        """
        try:
            result = run_command(
                ["podman", "exec", HAPROXY_CONTAINER, "haproxy", "-c", "-f", "/usr/local/etc/haproxy/haproxy.cfg"],
                timeout=SUBPROCESS_TIMEOUT,
            )
            if result.returncode == 0:
                return "Configuration is valid"
            return f"Configuration errors:\n{result.stderr}"
        except TimeoutError:
            return f"Error: Command timed out after {SUBPROCESS_TIMEOUT} seconds"
        except FileNotFoundError:
            return "Error: ssh/podman command not found"
        except OSError as e:
            return f"Error: OS error: {e}"

    @mcp.tool()
    def haproxy_save_state() -> str:
        """Save current server state to disk atomically.

        Returns:
            Success message or error description
        """
        try:
            state = haproxy_cmd("show servers state")
            atomic_write_file(STATE_FILE, state)
            return "Server state saved"
        except HaproxyError as e:
            return f"Error: {e}"
        except IOError as e:
            return f"Error: {e}"

    @mcp.tool()
    def haproxy_restore_state() -> str:
        """Restore server state from disk.

        Reads server configuration from servers.json and restores to HAProxy.

        Returns:
            Summary of restored servers or error description
        """
        try:
            restored = restore_servers_from_config()
            if restored == 0:
                return "No servers to restore"
            return f"Server state restored ({restored} servers)"
        except HaproxyError as e:
            return f"Error: {e}"
        except (OSError, ValueError) as e:
            return f"Error: {e}"