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,
    StateField,
    StatField,
    STATE_MIN_COLUMNS,
)
from ..exceptions import HaproxyError
from ..validation import validate_domain, validate_ip, validate_backend_name
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,
)


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


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)
        """
        if not validate_domain(domain):
            return "Error: Invalid domain format"

        try:
            backend, _ = get_backend_and_prefix(domain)
            servers = []
            state = haproxy_cmd("show servers state")

            for line in state.split("\n"):
                parts = line.split()
                if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.BE_NAME] == backend:
                    addr = parts[StateField.SRV_ADDR]
                    status = "active" if addr != "0.0.0.0" else "disabled"
                    servers.append(
                        f"• {parts[StateField.SRV_NAME]}: {addr}:{parts[StateField.SRV_PORT]} ({status})"
                    )

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

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

    @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)
        """
        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 (1 <= http_port <= 65535):
            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")
                used_slots: set[int] = set()
                for line in state.split("\n"):
                    parts = line.split()
                    if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.BE_NAME] == backend:
                        if parts[StateField.SRV_ADDR] != "0.0.0.0":
                            # Extract slot number from server name (e.g., pool_1_3 -> 3)
                            server_name = parts[StateField.SRV_NAME]
                            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}"

    @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"}]')
        """
        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 (1 <= http_port <= 65535):
                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 = []

        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)
                    added.append(f"slot {slot}: {ip}:{http_port}")
                except HaproxyError as e:
                    failed_slots.append(slot)
                    errors.append(f"slot {slot}: {e}")
        except Exception as e:
            # Rollback all saved configs on unexpected error
            for server_config in validated_servers:
                remove_server_from_config(domain, server_config["slot"])
            return f"Error: {e}"

        # Rollback failed slots from config
        for slot in failed_slots:
            remove_server_from_config(domain, slot)

        # 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"

    @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)
        """
        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}"
                # Batch both commands in single TCP connection
                haproxy_cmd_batch([
                    f"set server {backend}/{server} state maint",
                    f"set server {backend}/{server} addr 0.0.0.0 port 0"
                ])
                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}"

    @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")
        """
        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}"

        changed = []
        errors = []

        for line in servers_state.split("\n"):
            parts = line.split()
            if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.BE_NAME] == backend:
                server_name = parts[StateField.SRV_NAME]
                addr = parts[StateField.SRV_ADDR]
                # Only change state for configured servers (not 0.0.0.0)
                if 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

    @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)
        """
        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"

    @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")
        """
        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}"

    @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)
        """
        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}"