haproxy-mcp/haproxy_mcp/tools/domains.py

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

import fcntl
from typing import Annotated

from pydantic import Field

from ..config import (
    MAP_FILE,
    MAP_FILE_CONTAINER,
    POOL_COUNT,
    MAX_SLOTS,
    StateField,
    STATE_MIN_COLUMNS,
    logger,
)
from ..exceptions import HaproxyError
from ..validation import validate_domain, validate_ip
from ..haproxy_client import haproxy_cmd
from ..file_ops import (
    get_map_contents,
    save_map_file,
    get_domain_backend,
    is_legacy_backend,
    add_server_to_config,
    remove_server_from_config,
    remove_domain_from_config,
)


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

    @mcp.tool()
    def haproxy_list_domains(
        include_wildcards: Annotated[bool, Field(default=False, description="Include wildcard entries (.example.com). Default: False")]
    ) -> str:
        """List all configured domains with their backend servers."""
        try:
            domains = []
            state = haproxy_cmd("show servers state")

            # Build server map from HAProxy state
            server_map: dict[str, list] = {}
            for line in state.split("\n"):
                parts = line.split()
                if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.SRV_ADDR] != "0.0.0.0":
                    backend = parts[StateField.BE_NAME]
                    if backend not in server_map:
                        server_map[backend] = []
                    server_map[backend].append(
                        f"{parts[StateField.SRV_NAME]}={parts[StateField.SRV_ADDR]}:{parts[StateField.SRV_PORT]}"
                    )

            # Read from domains.map
            seen_domains: set[str] = set()
            for domain, backend in get_map_contents():
                # Skip wildcard entries unless explicitly requested
                if domain.startswith(".") and not include_wildcards:
                    continue
                if domain in seen_domains:
                    continue
                seen_domains.add(domain)
                servers = server_map.get(backend, ["(none)"])
                if domain.startswith("."):
                    backend_type = "wildcard"
                elif backend.startswith("pool_"):
                    backend_type = "pool"
                else:
                    backend_type = "static"
                domains.append(f"• {domain} -> {backend} ({backend_type}): {', '.join(servers)}")

            return "\n".join(domains) if domains else "No domains configured"
        except HaproxyError as e:
            return f"Error: {e}"

    @mcp.tool()
    def haproxy_add_domain(
        domain: Annotated[str, Field(description="Domain name to add (e.g., api.example.com, example.com)")],
        ip: Annotated[str, Field(default="", description="Optional: Initial server IP. If provided, adds server to slot 1")],
        http_port: Annotated[int, Field(default=80, description="HTTP port for backend server (default: 80)")]
    ) -> str:
        """Add a new domain to HAProxy (no reload required).

        Creates domain→pool mapping. Use haproxy_add_server to add more servers later.

        Example: haproxy_add_domain("api.example.com", ip="10.0.0.1", http_port=8080)
        """
        # Validate inputs
        if domain.startswith("."):
            return "Error: Domain cannot start with '.' (wildcard entries are added automatically)"
        if not validate_domain(domain):
            return "Error: Invalid domain format"
        if not validate_ip(ip, allow_empty=True):
            return "Error: Invalid IP address format"
        if not (1 <= http_port <= 65535):
            return "Error: Port must be between 1 and 65535"

        # Use file locking for the entire pool allocation operation
        lock_path = f"{MAP_FILE}.lock"
        with open(lock_path, 'w') as lock_file:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            try:
                # Read map contents once for both existence check and pool lookup
                entries = get_map_contents()

                # Check if domain already exists (using cached entries)
                for domain_entry, backend in entries:
                    if domain_entry == domain:
                        return f"Error: Domain {domain} already exists (mapped to {backend})"

                # Find available pool (using cached entries)
                used_pools: set[str] = set()
                for _, backend in entries:
                    if backend.startswith("pool_"):
                        used_pools.add(backend)

                pool = None
                for i in range(1, POOL_COUNT + 1):
                    pool_name = f"pool_{i}"
                    if pool_name not in used_pools:
                        pool = pool_name
                        break
                if not pool:
                    return f"Error: All {POOL_COUNT} pool backends are in use"

                try:
                    # Save to disk first (atomic write for persistence)
                    entries.append((domain, pool))
                    entries.append((f".{domain}", pool))
                    try:
                        save_map_file(entries)
                    except IOError as e:
                        return f"Error: Failed to save map file: {e}"

                    # Then update HAProxy map via Runtime API
                    try:
                        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} {domain} {pool}")
                        haproxy_cmd(f"add map {MAP_FILE_CONTAINER} .{domain} {pool}")
                    except HaproxyError as e:
                        # Rollback: remove the domain we just added from entries and re-save
                        rollback_entries = [(d, b) for d, b in entries if d != domain and d != f".{domain}"]
                        try:
                            save_map_file(rollback_entries)
                        except IOError:
                            logger.error("Failed to rollback map file after HAProxy error")
                        return f"Error: Failed to update HAProxy map: {e}"

                    # If IP provided, add server to slot 1
                    if ip:
                        # Save server config to disk first
                        add_server_to_config(domain, 1, ip, http_port)

                        try:
                            server = f"{pool}_1"
                            haproxy_cmd(f"set server {pool}/{server} addr {ip} port {http_port}")
                            haproxy_cmd(f"set server {pool}/{server} state ready")
                        except HaproxyError as e:
                            # Rollback server config on failure
                            remove_server_from_config(domain, 1)
                            return f"Domain {domain} added to {pool} but server config failed: {e}"

                        return f"Domain {domain} added to {pool} with server {ip}:{http_port}"

                    return f"Domain {domain} added to {pool} (no servers configured)"

                except HaproxyError as e:
                    return f"Error: {e}"
            finally:
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)

    @mcp.tool()
    def haproxy_remove_domain(
        domain: Annotated[str, Field(description="Domain name to remove (e.g., api.example.com)")]
    ) -> str:
        """Remove a domain from HAProxy (no reload required)."""
        if not validate_domain(domain):
            return "Error: Invalid domain format"

        # Look up the domain in the map
        backend = get_domain_backend(domain)
        if not backend:
            return f"Error: Domain {domain} not found"

        # Check if this is a legacy backend (not a pool)
        if is_legacy_backend(backend):
            return f"Error: Cannot remove legacy domain {domain} (uses static backend {backend})"

        try:
            # Save to disk first (atomic write for persistence)
            entries = get_map_contents()
            new_entries = [(d, b) for d, b in entries if d != domain and d != f".{domain}"]
            save_map_file(new_entries)

            # Remove from persistent server config
            remove_domain_from_config(domain)

            # Clear map entries via Runtime API (immediate effect)
            haproxy_cmd(f"del map {MAP_FILE_CONTAINER} {domain}")
            try:
                haproxy_cmd(f"del map {MAP_FILE_CONTAINER} .{domain}")
            except HaproxyError as e:
                logger.warning("Failed to remove wildcard entry for %s: %s", domain, e)

            # Disable all servers in the pool (reset to 0.0.0.0:0)
            for slot in range(1, MAX_SLOTS + 1):
                server = f"{backend}_{slot}"
                try:
                    haproxy_cmd(f"set server {backend}/{server} state maint")
                    haproxy_cmd(f"set server {backend}/{server} addr 0.0.0.0 port 0")
                except HaproxyError as e:
                    logger.warning(
                        "Failed to clear server %s/%s for domain %s: %s",
                        backend, server, domain, e
                    )
                    # Continue with remaining cleanup

            return f"Domain {domain} removed from {backend}"

        except IOError as e:
            return f"Error: Failed to update map file: {e}"
        except HaproxyError as e:
            return f"Error: {e}"