haproxy-mcp/haproxy_mcp/tools/servers.py

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

import json
import time
from typing import Annotated

from pydantic import Field

from ..config import (
    MAX_SLOTS,
    MAX_BULK_SERVERS,
    MAX_SERVERS_JSON_SIZE,
    StatField,
    logger,
)
from ..exceptions import HaproxyError
from ..validation import validate_domain, validate_ip, validate_backend_name, validate_port_int
from ..haproxy_client import haproxy_cmd, haproxy_cmd_checked, haproxy_cmd_batch
from ..file_ops import (
    get_backend_and_prefix,
    load_servers_config,
    add_server_to_config,
    remove_server_from_config,
)
from ..utils import parse_servers_state, disable_server_slot


def configure_server_slot(backend: str, server_prefix: str, slot: int, ip: str, http_port: int) -> str:
    """Configure a server slot in HAProxy.

    Args:
        backend: Backend name (e.g., 'pool_1')
        server_prefix: Server name prefix (e.g., 'pool_1')
        slot: Slot number (1-10)
        ip: Server IP address
        http_port: HTTP port

    Returns:
        Server name that was configured

    Raises:
        HaproxyError: If HAProxy command fails
    """
    server = f"{server_prefix}_{slot}"
    # Batch both commands in single TCP connection
    haproxy_cmd_batch([
        f"set server {backend}/{server} addr {ip} port {http_port}",
        f"set server {backend}/{server} state ready"
    ])
    return server


# =============================================================================
# Implementation functions (module-level)
# =============================================================================


def _haproxy_list_servers_impl(domain: str) -> str:
    """Implementation of haproxy_list_servers."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"

    try:
        backend, _ = get_backend_and_prefix(domain)
        state = haproxy_cmd("show servers state")
        parsed_state = parse_servers_state(state)
        backend_servers = parsed_state.get(backend, {})

        if not backend_servers:
            return f"Backend {backend} not found"

        servers = []
        for server_name, srv_info in backend_servers.items():
            addr = srv_info["addr"]
            port = srv_info["port"]
            status = "active" if addr != "0.0.0.0" else "disabled"
            servers.append(f"• {server_name}: {addr}:{port} ({status})")

        return f"Servers for {domain} ({backend}):\n" + "\n".join(servers)
    except (HaproxyError, ValueError) as e:
        return f"Error: {e}"


def _haproxy_add_server_impl(domain: str, slot: int, ip: str, http_port: int) -> str:
    """Implementation of haproxy_add_server."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"
    if not ip:
        return "Error: IP address is required"
    if not validate_ip(ip):
        return "Error: Invalid IP address format"
    if not validate_port_int(http_port):
        return "Error: Port must be between 1 and 65535"

    try:
        backend, server_prefix = get_backend_and_prefix(domain)

        # Auto-select slot if slot <= 0
        if slot <= 0:
            state = haproxy_cmd("show servers state")
            parsed_state = parse_servers_state(state)
            backend_servers = parsed_state.get(backend, {})
            used_slots: set[int] = set()
            for server_name, srv_info in backend_servers.items():
                if srv_info["addr"] != "0.0.0.0":
                    # Extract slot number from server name (e.g., pool_1_3 -> 3)
                    try:
                        used_slots.add(int(server_name.rsplit("_", 1)[1]))
                    except (ValueError, IndexError):
                        pass
            for s in range(1, MAX_SLOTS + 1):
                if s not in used_slots:
                    slot = s
                    break
            else:
                return f"Error: No available slots (all {MAX_SLOTS} slots in use)"
        elif not (1 <= slot <= MAX_SLOTS):
            return f"Error: Slot must be between 1 and {MAX_SLOTS}, or 0/-1 for auto-select"

        # Save to persistent config FIRST (disk-first pattern)
        add_server_to_config(domain, slot, ip, http_port)

        try:
            server = configure_server_slot(backend, server_prefix, slot, ip, http_port)
            return f"Added to {domain} ({backend}) slot {slot}:\n{server} → {ip}:{http_port}"
        except HaproxyError as e:
            # Rollback config on HAProxy failure
            remove_server_from_config(domain, slot)
            return f"Error: {e}"
    except (ValueError, IOError) as e:
        return f"Error: {e}"


def _haproxy_add_servers_impl(domain: str, servers: str) -> str:
    """Implementation of haproxy_add_servers."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"

    # Check JSON size before parsing
    if len(servers) > MAX_SERVERS_JSON_SIZE:
        return f"Error: servers JSON exceeds maximum size ({MAX_SERVERS_JSON_SIZE} bytes)"

    # Parse JSON array
    try:
        server_list = json.loads(servers)
    except json.JSONDecodeError as e:
        return f"Error: Invalid JSON - {e}"

    if not isinstance(server_list, list):
        return "Error: servers must be a JSON array"

    if not server_list:
        return "Error: servers array is empty"

    if len(server_list) > MAX_BULK_SERVERS:
        return f"Error: Cannot add more than {MAX_BULK_SERVERS} servers at once"

    # Validate all servers first before adding any
    validated_servers = []
    validation_errors = []

    for i, srv in enumerate(server_list):
        if not isinstance(srv, dict):
            validation_errors.append(f"Server {i+1}: must be an object")
            continue

        # Extract and validate slot
        slot = srv.get("slot")
        if slot is None:
            validation_errors.append(f"Server {i+1}: missing 'slot' field")
            continue
        try:
            slot = int(slot)
        except (ValueError, TypeError):
            validation_errors.append(f"Server {i+1}: slot must be an integer")
            continue
        if not (1 <= slot <= MAX_SLOTS):
            validation_errors.append(f"Server {i+1}: slot must be between 1 and {MAX_SLOTS}")
            continue

        # Extract and validate IP
        ip = srv.get("ip")
        if not ip:
            validation_errors.append(f"Server {i+1}: missing 'ip' field")
            continue
        if not validate_ip(ip):
            validation_errors.append(f"Server {i+1}: invalid IP address '{ip}'")
            continue

        # Extract and validate port
        http_port = srv.get("http_port", 80)
        try:
            http_port = int(http_port)
        except (ValueError, TypeError):
            validation_errors.append(f"Server {i+1}: http_port must be an integer")
            continue
        if not validate_port_int(http_port):
            validation_errors.append(f"Server {i+1}: port must be between 1 and 65535")
            continue

        validated_servers.append({"slot": slot, "ip": ip, "http_port": http_port})

    # Return validation errors if any
    if validation_errors:
        return "Validation errors:\n" + "\n".join(f"  • {e}" for e in validation_errors)

    # Check for duplicate slots
    slots = [s["slot"] for s in validated_servers]
    if len(slots) != len(set(slots)):
        return "Error: Duplicate slot numbers in servers array"

    # Get backend info
    try:
        backend, server_prefix = get_backend_and_prefix(domain)
    except ValueError as e:
        return f"Error: {e}"

    # Save ALL servers to config FIRST (disk-first pattern)
    for server_config in validated_servers:
        slot = server_config["slot"]
        ip = server_config["ip"]
        http_port = server_config["http_port"]
        add_server_to_config(domain, slot, ip, http_port)

    # Then update HAProxy
    added = []
    errors = []
    failed_slots = []
    successfully_added_slots = []

    try:
        for server_config in validated_servers:
            slot = server_config["slot"]
            ip = server_config["ip"]
            http_port = server_config["http_port"]
            try:
                configure_server_slot(backend, server_prefix, slot, ip, http_port)
                successfully_added_slots.append(slot)
                added.append(f"slot {slot}: {ip}:{http_port}")
            except HaproxyError as e:
                failed_slots.append(slot)
                errors.append(f"slot {slot}: {e}")
    except (IOError, OSError, ValueError) as e:
        # Rollback only successfully added configs on unexpected error
        logger.error("Unexpected error during bulk server add for %s: %s", domain, e)
        for slot in successfully_added_slots:
            try:
                remove_server_from_config(domain, slot)
            except (IOError, OSError, ValueError) as rollback_error:
                logger.error(
                    "Failed to rollback server config for %s slot %d: %s",
                    domain, slot, rollback_error
                )
        # Also rollback configs that weren't yet processed
        for server_config in validated_servers:
            slot = server_config["slot"]
            if slot not in successfully_added_slots:
                try:
                    remove_server_from_config(domain, slot)
                except (IOError, OSError, ValueError) as rollback_error:
                    logger.error(
                        "Failed to rollback server config for %s slot %d: %s",
                        domain, slot, rollback_error
                    )
        return f"Error: {e}"

    # Rollback failed slots from config
    for slot in failed_slots:
        try:
            remove_server_from_config(domain, slot)
        except (IOError, OSError, ValueError) as rollback_error:
            logger.error(
                "Failed to rollback server config for %s slot %d: %s",
                domain, slot, rollback_error
            )

    # Build result message
    result_parts = []
    if added:
        result_parts.append(f"Added {len(added)} servers to {domain} ({backend}):")
        result_parts.extend(f"  • {s}" for s in added)
    if errors:
        result_parts.append(f"Failed to add {len(errors)} servers:")
        result_parts.extend(f"  • {e}" for e in errors)

    return "\n".join(result_parts) if result_parts else "No servers added"


def _haproxy_remove_server_impl(domain: str, slot: int) -> str:
    """Implementation of haproxy_remove_server."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"
    if not (1 <= slot <= MAX_SLOTS):
        return f"Error: Slot must be between 1 and {MAX_SLOTS}"

    try:
        backend, server_prefix = get_backend_and_prefix(domain)

        # Get current server info for potential rollback
        config = load_servers_config()
        old_config = config.get(domain, {}).get(str(slot), {})

        # Remove from persistent config FIRST (disk-first pattern)
        remove_server_from_config(domain, slot)

        try:
            # HTTP only - single server per slot
            server = f"{server_prefix}_{slot}"
            disable_server_slot(backend, server)
            return f"Removed server at slot {slot} from {domain} ({backend})"
        except HaproxyError as e:
            # Rollback: re-add config if HAProxy command failed
            if old_config:
                add_server_to_config(domain, slot, old_config.get("ip", ""), old_config.get("http_port", 80))
            return f"Error: {e}"
    except (ValueError, IOError) as e:
        return f"Error: {e}"


def _haproxy_set_domain_state_impl(domain: str, state: str) -> str:
    """Implementation of haproxy_set_domain_state."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"
    if state not in ["ready", "drain", "maint"]:
        return "Error: State must be 'ready', 'drain', or 'maint'"

    try:
        backend, _ = get_backend_and_prefix(domain)
    except ValueError as e:
        return f"Error: {e}"

    # Get active servers for this domain
    try:
        servers_state = haproxy_cmd("show servers state")
    except HaproxyError as e:
        return f"Error: {e}"

    parsed_state = parse_servers_state(servers_state)
    backend_servers = parsed_state.get(backend, {})

    changed = []
    errors = []

    for server_name, srv_info in backend_servers.items():
        # Only change state for configured servers (not 0.0.0.0)
        if srv_info["addr"] != "0.0.0.0":
            try:
                haproxy_cmd_checked(f"set server {backend}/{server_name} state {state}")
                changed.append(server_name)
            except HaproxyError as e:
                errors.append(f"{server_name}: {e}")

    if not changed and not errors:
        return f"No active servers found for {domain}"

    result = f"Set {len(changed)} servers to '{state}' for {domain}"
    if changed:
        result += ":\n" + "\n".join(f"  • {s}" for s in changed)
    if errors:
        result += f"\n\nErrors ({len(errors)}):\n" + "\n".join(f"  • {e}" for e in errors)

    return result


def _haproxy_wait_drain_impl(domain: str, timeout: int) -> str:
    """Implementation of haproxy_wait_drain."""
    if not validate_domain(domain):
        return "Error: Invalid domain format"
    if not (1 <= timeout <= 300):
        return "Error: Timeout must be between 1 and 300 seconds"

    try:
        backend, _ = get_backend_and_prefix(domain)
    except ValueError as e:
        return f"Error: {e}"

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            stats = haproxy_cmd("show stat")
            total_connections = 0
            for line in stats.split("\n"):
                parts = line.split(",")
                if len(parts) > StatField.SCUR and parts[0] == backend and parts[1] not in ["FRONTEND", "BACKEND", ""]:
                    try:
                        scur = int(parts[StatField.SCUR]) if parts[StatField.SCUR] else 0
                        total_connections += scur
                    except ValueError:
                        pass

            if total_connections == 0:
                elapsed = int(time.time() - start_time)
                return f"All connections drained for {domain} ({elapsed}s)"

            time.sleep(1)
        except HaproxyError as e:
            return f"Error checking connections: {e}"

    return f"Timeout: Connections still active after {timeout}s"


def _haproxy_set_server_state_impl(backend: str, server: str, state: str) -> str:
    """Implementation of haproxy_set_server_state."""
    if not validate_backend_name(backend):
        return "Error: Invalid backend name (use alphanumeric, underscore, hyphen only)"
    if not validate_backend_name(server):
        return "Error: Invalid server name (use alphanumeric, underscore, hyphen only)"
    if state not in ["ready", "drain", "maint"]:
        return "Error: state must be 'ready', 'drain', or 'maint'"
    try:
        haproxy_cmd_checked(f"set server {backend}/{server} state {state}")
        return f"Server {backend}/{server} set to {state}"
    except HaproxyError as e:
        return f"Error: {e}"


def _haproxy_set_server_weight_impl(backend: str, server: str, weight: int) -> str:
    """Implementation of haproxy_set_server_weight."""
    if not validate_backend_name(backend):
        return "Error: Invalid backend name (use alphanumeric, underscore, hyphen only)"
    if not validate_backend_name(server):
        return "Error: Invalid server name (use alphanumeric, underscore, hyphen only)"
    if not (0 <= weight <= 256):
        return "Error: weight must be between 0 and 256"
    try:
        haproxy_cmd_checked(f"set server {backend}/{server} weight {weight}")
        return f"Server {backend}/{server} weight set to {weight}"
    except HaproxyError as e:
        return f"Error: {e}"


# =============================================================================
# MCP Tool Registration
# =============================================================================


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

    @mcp.tool()
    def haproxy_list_servers(
        domain: Annotated[str, Field(description="Domain name to list servers for (e.g., api.example.com)")]
    ) -> str:
        """List all servers for a domain with slot numbers, addresses, and status (UP/DOWN/MAINT).

        Example:
            haproxy_list_servers("api.example.com")
            # Output: pool_1_1: 10.0.0.1:8080 (UP)
            #         pool_1_2: 10.0.0.2:8080 (UP)
        """
        return _haproxy_list_servers_impl(domain)

    @mcp.tool()
    def haproxy_add_server(
        domain: Annotated[str, Field(description="Domain name to add server to (e.g., api.example.com)")],
        slot: Annotated[int, Field(description="Server slot number 1-10, or 0 for auto-select next available slot")],
        ip: Annotated[str, Field(description="Server IP address (IPv4 like 10.0.0.1 or IPv6 like 2001:db8::1)")],
        http_port: Annotated[int, Field(default=80, description="HTTP port for backend connection (default: 80)")]
    ) -> str:
        """Add a server to a domain's backend pool for load balancing.

        Each domain can have up to 10 servers (slots 1-10). HAProxy distributes traffic
        across all configured servers using round-robin.

        Example: haproxy_add_server("api.example.com", slot=1, ip="10.0.0.1", http_port=8080)
        """
        return _haproxy_add_server_impl(domain, slot, ip, http_port)

    @mcp.tool()
    def haproxy_add_servers(
        domain: Annotated[str, Field(description="Domain name to add servers to (e.g., api.example.com)")],
        servers: Annotated[str, Field(description='JSON array of servers. Each object: {"slot": 1-10, "ip": "10.0.0.1", "http_port": 80}. Example: \'[{"slot":1,"ip":"10.0.0.1"},{"slot":2,"ip":"10.0.0.2"}]\'')]
    ) -> str:
        """Add multiple servers to a domain's backend at once (bulk operation).

        Example: haproxy_add_servers("api.example.com", '[{"slot":1,"ip":"10.0.0.1"},{"slot":2,"ip":"10.0.0.2"}]')
        """
        return _haproxy_add_servers_impl(domain, servers)

    @mcp.tool()
    def haproxy_remove_server(
        domain: Annotated[str, Field(description="Domain name to remove server from (e.g., api.example.com)")],
        slot: Annotated[int, Field(description="Server slot number to remove (1-10)")]
    ) -> str:
        """Remove a server from a domain's backend at specified slot.

        Example: haproxy_remove_server("api.example.com", slot=2)
        """
        return _haproxy_remove_server_impl(domain, slot)

    @mcp.tool()
    def haproxy_set_domain_state(
        domain: Annotated[str, Field(description="Domain name (e.g., api.example.com)")],
        state: Annotated[str, Field(description="Target state: 'ready' (normal), 'drain' (stop new connections), or 'maint' (maintenance)")]
    ) -> str:
        """Set state for all servers of a domain at once.

        Example: haproxy_set_domain_state("api.example.com", state="drain")

        Example:
            # Put all servers in maintenance for deployment
            haproxy_set_domain_state("api.example.com", "maint")

            # Re-enable all servers after deployment
            haproxy_set_domain_state("api.example.com", "ready")
        """
        return _haproxy_set_domain_state_impl(domain, state)

    @mcp.tool()
    def haproxy_wait_drain(
        domain: Annotated[str, Field(description="Domain name to wait for (e.g., api.example.com)")],
        timeout: Annotated[int, Field(default=30, description="Maximum seconds to wait (default: 30, max: 300)")]
    ) -> str:
        """Wait for all active connections to drain from a domain's servers.

        Use after setting servers to 'drain' state before maintenance.

        Example: haproxy_wait_drain("api.example.com", timeout=60)
        """
        return _haproxy_wait_drain_impl(domain, timeout)

    @mcp.tool()
    def haproxy_set_server_state(
        backend: Annotated[str, Field(description="Backend name (e.g., 'pool_1')")],
        server: Annotated[str, Field(description="Server name (e.g., 'pool_1_1')")],
        state: Annotated[str, Field(description="'ready' (enable), 'drain' (graceful shutdown), or 'maint' (maintenance)")]
    ) -> str:
        """Set server state for maintenance or traffic control.

        Example: haproxy_set_server_state("pool_1", "pool_1_2", "maint")
        """
        return _haproxy_set_server_state_impl(backend, server, state)

    @mcp.tool()
    def haproxy_set_server_weight(
        backend: Annotated[str, Field(description="Backend name (e.g., 'pool_1')")],
        server: Annotated[str, Field(description="Server name (e.g., 'pool_1_1')")],
        weight: Annotated[int, Field(description="Weight 0-256 (higher = more traffic, 0 = disabled)")]
    ) -> str:
        """Set server weight for load balancing ratio control.

        Example: haproxy_set_server_weight("pool_1", "pool_1_1", weight=2)
        """
        return _haproxy_set_server_weight_impl(backend, server, weight)