haproxy-mcp/haproxy_mcp/tools/configuration.py

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

import fcntl
import subprocess
import time

from ..config import (
    STATE_FILE,
    HAPROXY_CONTAINER,
    SUBPROCESS_TIMEOUT,
    STARTUP_RETRY_COUNT,
    StateField,
    STATE_MIN_COLUMNS,
    logger,
)
from ..exceptions import HaproxyError
from ..validation import validate_ip, validate_port, validate_backend_name
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,
)


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 in single batch
    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 from config file 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

    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)


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

        # Restore servers from config after reload
        try:
            restored = restore_servers_from_config()
            return f"HAProxy configuration reloaded successfully ({restored} servers restored)"
        except Exception 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 = subprocess.run(
                ["podman", "exec", HAPROXY_CONTAINER, "haproxy", "-c", "-f", "/usr/local/etc/haproxy/haproxy.cfg"],
                capture_output=True, text=True, timeout=SUBPROCESS_TIMEOUT
            )
            if result.returncode == 0:
                return "Configuration is valid"
            return f"Configuration errors:\n{result.stderr}"
        except subprocess.TimeoutExpired:
            return f"Error: Command timed out after {SUBPROCESS_TIMEOUT} seconds"
        except FileNotFoundError:
            return "Error: 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.

        Uses batched commands for efficiency.

        Returns:
            Summary of restored servers or error description
        """
        try:
            with open(STATE_FILE, "r", encoding="utf-8") as f:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_SH)
                except OSError:
                    pass  # Continue without lock if not supported
                try:
                    state = f.read()
                finally:
                    try:
                        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
                    except OSError:
                        pass

            # Build batch of all commands
            commands: list[str] = []
            server_info_list: list[tuple[str, str]] = []
            skipped = 0

            for line in state.split("\n"):
                parts = line.split()
                if len(parts) >= STATE_MIN_COLUMNS and not line.startswith("#"):
                    backend = parts[StateField.BE_NAME]
                    server = parts[StateField.SRV_NAME]
                    addr = parts[StateField.SRV_ADDR]
                    port = parts[StateField.SRV_PORT]

                    # Skip disabled servers
                    if addr == "0.0.0.0":
                        continue

                    # Validate names from state file to prevent injection
                    if not validate_backend_name(backend) or not validate_backend_name(server):
                        skipped += 1
                        continue

                    # Validate IP and port
                    if not validate_ip(addr) or not validate_port(port):
                        skipped += 1
                        continue

                    commands.append(f"set server {backend}/{server} addr {addr} port {port}")
                    commands.append(f"set server {backend}/{server} state ready")
                    server_info_list.append((backend, server))

            if not commands:
                result = "No servers to restore"
                if skipped:
                    result += f", {skipped} entries skipped due to validation"
                return result

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

            result = f"Server state restored ({restored} servers)"
            if skipped:
                result += f", {skipped} entries skipped due to validation"
            return result
        except FileNotFoundError:
            return "Error: No saved state found"
        except HaproxyError as e:
            return f"Error: {e}"