haproxy-mcp/haproxy_mcp/file_ops.py

"""File I/O operations for HAProxy MCP Server."""

import fcntl
import json
import os
import tempfile
from contextlib import contextmanager
from typing import Any, Generator, Optional

from .config import (
    MAP_FILE,
    WILDCARDS_MAP_FILE,
    SERVERS_FILE,
    CERTS_FILE,
    logger,
)
from .validation import domain_to_backend


@contextmanager
def file_lock(lock_path: str) -> Generator[None, None, None]:
    """Acquire exclusive file lock for atomic operations.

    This context manager provides a consistent locking mechanism for
    read-modify-write operations on configuration files to prevent
    race conditions during concurrent access.

    Args:
        lock_path: Path to the lock file (typically config_file.lock)

    Yields:
        None - the lock is held for the duration of the context

    Example:
        with file_lock("/path/to/config.json.lock"):
            config = load_config()
            config["key"] = "value"
            save_config(config)
    """
    with open(lock_path, 'w') as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)


def atomic_write_file(file_path: str, content: str) -> None:
    """Write content to file atomically using temp file + rename.

    Args:
        file_path: Target file path
        content: Content to write

    Raises:
        IOError: If write fails
    """
    dir_path = os.path.dirname(file_path)
    fd = None
    temp_path = None
    try:
        fd, temp_path = tempfile.mkstemp(dir=dir_path, prefix='.tmp.')
        with os.fdopen(fd, 'w', encoding='utf-8') as f:
            fd = None  # fd is now owned by the file object
            f.write(content)
        os.rename(temp_path, file_path)
        temp_path = None  # Rename succeeded
    except OSError as e:
        raise IOError(f"Failed to write {file_path}: {e}") from e
    finally:
        if fd is not None:
            try:
                os.close(fd)
            except OSError:
                pass
        if temp_path is not None:
            try:
                os.unlink(temp_path)
            except OSError:
                pass


def _read_map_file(file_path: str) -> list[tuple[str, str]]:
    """Read a single map file and return list of (domain, backend) tuples.

    Args:
        file_path: Path to the map file

    Returns:
        List of (domain, backend) tuples from the map file
    """
    entries = []
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            try:
                fcntl.flock(f.fileno(), fcntl.LOCK_SH)
            except OSError:
                pass  # Continue without lock if not supported
            try:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#"):
                        continue
                    parts = line.split()
                    if len(parts) >= 2:
                        entries.append((parts[0], parts[1]))
            finally:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
                except OSError:
                    pass
    except FileNotFoundError:
        pass
    return entries


def get_map_contents() -> list[tuple[str, str]]:
    """Read both domains.map and wildcards.map and return combined entries.

    Returns:
        List of (domain, backend) tuples from both map files
    """
    # Read exact domains
    entries = _read_map_file(MAP_FILE)
    # Read wildcards and append
    entries.extend(_read_map_file(WILDCARDS_MAP_FILE))
    return entries


def split_domain_entries(entries: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]:
    """Split entries into exact domains and wildcards.

    Args:
        entries: List of (domain, backend) tuples

    Returns:
        Tuple of (exact_entries, wildcard_entries)
    """
    exact = []
    wildcards = []
    for domain, backend in entries:
        if domain.startswith("."):
            wildcards.append((domain, backend))
        else:
            exact.append((domain, backend))
    return exact, wildcards


def save_map_file(entries: list[tuple[str, str]]) -> None:
    """Save domain-to-backend entries using 2-stage map routing architecture.

    This function implements HAProxy's 2-stage domain routing for optimal
    performance. Entries are automatically split into two separate map files
    based on whether they are exact domains or wildcard patterns.

    2-Stage Routing Architecture:
        Stage 1 - Exact Match (domains.map):
            - HAProxy directive: map_str(req.hdr(host),"/path/domains.map")
            - Data structure: ebtree (elastic binary tree)
            - Lookup complexity: O(log n)
            - Use case: Exact domain matches (e.g., "api.example.com")

        Stage 2 - Wildcard Match (wildcards.map):
            - HAProxy directive: map_dom(req.hdr(host),"/path/wildcards.map")
            - Data structure: Linear suffix search
            - Lookup complexity: O(n) where n = number of wildcard entries
            - Use case: Wildcard domains (e.g., ".example.com" matches *.example.com)
            - Typically small set, so O(n) is acceptable

    Performance Characteristics:
        - 1000 exact domains: ~10 comparisons (log2(1000) approx 10)
        - 10 wildcard entries: 10 suffix comparisons (acceptable)
        - By separating exact and wildcard entries, we avoid O(n) lookup
          for the common case (exact domain match)

    HAProxy Configuration Example:
        use_backend %[req.hdr(host),lower,map_str(/etc/haproxy/domains.map)]
            if { req.hdr(host),lower,map_str(/etc/haproxy/domains.map) -m found }
        use_backend %[req.hdr(host),lower,map_dom(/etc/haproxy/wildcards.map)]
            if { req.hdr(host),lower,map_dom(/etc/haproxy/wildcards.map) -m found }

    Args:
        entries: List of (domain, backend) tuples to write.
            - Exact domains: "api.example.com" -> written to domains.map
            - Wildcards: ".example.com" (matches *.example.com) -> written
              to wildcards.map

    Raises:
        IOError: If either map file cannot be written.

    File Formats:
        domains.map:
            # Exact Domain to Backend mapping (for map_str)
            api.example.com pool_1
            www.example.com pool_2

        wildcards.map:
            # Wildcard Domain to Backend mapping (for map_dom)
            .example.com pool_3    # Matches *.example.com
            .test.org pool_4       # Matches *.test.org

    Note:
        Both files are written atomically using temp file + rename to prevent
        corruption during concurrent access or system failures.
    """
    # Split into exact and wildcard entries
    exact_entries, wildcard_entries = split_domain_entries(entries)

    # Save exact domains (for map_str - fast O(log n) lookup)
    exact_lines = [
        "# Exact Domain to Backend mapping (for map_str)\n",
        "# Format: domain backend_name\n",
        "# Uses ebtree for O(log n) lookup performance\n\n",
    ]
    for domain, backend in sorted(exact_entries):
        exact_lines.append(f"{domain} {backend}\n")
    atomic_write_file(MAP_FILE, "".join(exact_lines))

    # Save wildcards (for map_dom - O(n) but small set)
    wildcard_lines = [
        "# Wildcard Domain to Backend mapping (for map_dom)\n",
        "# Format: .domain.com backend_name (matches *.domain.com)\n",
        "# Uses map_dom for suffix matching\n\n",
    ]
    for domain, backend in sorted(wildcard_entries):
        wildcard_lines.append(f"{domain} {backend}\n")
    atomic_write_file(WILDCARDS_MAP_FILE, "".join(wildcard_lines))


def get_domain_backend(domain: str) -> Optional[str]:
    """Look up the backend for a domain from domains.map.

    Args:
        domain: The domain to look up

    Returns:
        Backend name if found, None otherwise
    """
    for map_domain, backend in get_map_contents():
        if map_domain == domain:
            return backend
    return None


def is_legacy_backend(backend: str) -> bool:
    """Check if backend is a legacy static backend (not a dynamic pool).

    This function distinguishes between two backend naming conventions used
    in the HAProxy MCP system:

    Pool Backends (Dynamic):
        - Named: pool_1, pool_2, ..., pool_100
        - Pre-configured in haproxy.cfg with 10 server slots each
        - Domains are dynamically assigned to available pools via domains.map
        - Server slots configured at runtime via Runtime API
        - Allows zero-reload domain management

    Legacy Backends (Static):
        - Named: {domain}_backend (e.g., "api_example_com_backend")
        - Defined statically in haproxy.cfg
        - Requires HAProxy reload to add new backends
        - Used for domains that were configured before pool-based routing

    Args:
        backend: Backend name to check (e.g., "pool_5" or "api_example_com_backend").

    Returns:
        True if this is a legacy backend (does not start with "pool_"),
        False if it's a pool backend.

    Usage Scenarios:
        - When listing servers: Determines server naming convention
          (pool backends use pool_N_M, legacy use {domain}_M)
        - When adding servers: Determines which backend configuration
          approach to use
        - During migration: Helps identify domains that need migration
          from legacy to pool-based routing

    Examples:
        >>> is_legacy_backend("pool_5")
        False
        >>> is_legacy_backend("pool_100")
        False
        >>> is_legacy_backend("api_example_com_backend")
        True
        >>> is_legacy_backend("myservice_backend")
        True
    """
    return not backend.startswith("pool_")


def get_legacy_backend_name(domain: str) -> str:
    """Convert domain to legacy backend name format.

    Args:
        domain: Domain name

    Returns:
        Legacy backend name (e.g., 'api_example_com_backend')
    """
    return f"{domain_to_backend(domain)}_backend"


def get_backend_and_prefix(domain: str) -> tuple[str, str]:
    """Look up backend and determine server name prefix for a domain.

    Args:
        domain: The domain name to look up

    Returns:
        Tuple of (backend_name, server_prefix)

    Raises:
        ValueError: If domain cannot be mapped to a valid backend
    """
    backend = get_domain_backend(domain)
    if not backend:
        backend = get_legacy_backend_name(domain)

    if backend.startswith("pool_"):
        server_prefix = backend
    else:
        server_prefix = domain_to_backend(domain)

    return backend, server_prefix


def load_servers_config() -> dict[str, Any]:
    """Load servers configuration from JSON file with file locking.

    Returns:
        Dictionary with server configurations
    """
    try:
        with open(SERVERS_FILE, "r", encoding="utf-8") as f:
            try:
                fcntl.flock(f.fileno(), fcntl.LOCK_SH)
            except OSError:
                logger.debug("File locking not supported for %s", SERVERS_FILE)
            try:
                return json.load(f)
            finally:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
                except OSError:
                    pass
    except FileNotFoundError:
        return {}
    except json.JSONDecodeError as e:
        logger.warning("Corrupt config file %s: %s", SERVERS_FILE, e)
        return {}


def save_servers_config(config: dict[str, Any]) -> None:
    """Save servers configuration to JSON file atomically.

    Uses temp file + rename for atomic write to prevent race conditions.

    Args:
        config: Dictionary with server configurations
    """
    atomic_write_file(SERVERS_FILE, json.dumps(config, indent=2))


def add_server_to_config(domain: str, slot: int, ip: str, http_port: int) -> None:
    """Add server configuration to persistent storage with file locking.

    Args:
        domain: Domain name
        slot: Server slot (1 to MAX_SLOTS)
        ip: Server IP address
        http_port: HTTP port
    """
    with file_lock(f"{SERVERS_FILE}.lock"):
        config = load_servers_config()
        if domain not in config:
            config[domain] = {}
        config[domain][str(slot)] = {"ip": ip, "http_port": http_port}
        save_servers_config(config)


def remove_server_from_config(domain: str, slot: int) -> None:
    """Remove server configuration from persistent storage with file locking.

    Args:
        domain: Domain name
        slot: Server slot to remove
    """
    with file_lock(f"{SERVERS_FILE}.lock"):
        config = load_servers_config()
        if domain in config and str(slot) in config[domain]:
            del config[domain][str(slot)]
            if not config[domain]:
                del config[domain]
            save_servers_config(config)


def remove_domain_from_config(domain: str) -> None:
    """Remove domain from persistent config with file locking.

    Args:
        domain: Domain name to remove
    """
    with file_lock(f"{SERVERS_FILE}.lock"):
        config = load_servers_config()
        if domain in config:
            del config[domain]
            save_servers_config(config)


# Certificate configuration functions

def load_certs_config() -> list[str]:
    """Load certificate domain list from JSON file.

    Returns:
        List of domain names
    """
    try:
        with open(CERTS_FILE, "r", encoding="utf-8") as f:
            try:
                fcntl.flock(f.fileno(), fcntl.LOCK_SH)
            except OSError:
                pass
            try:
                data = json.load(f)
                return data.get("domains", [])
            finally:
                try:
                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
                except OSError:
                    pass
    except FileNotFoundError:
        return []
    except json.JSONDecodeError as e:
        logger.warning("Corrupt certificates config %s: %s", CERTS_FILE, e)
        return []


def save_certs_config(domains: list[str]) -> None:
    """Save certificate domain list to JSON file atomically.

    Args:
        domains: List of domain names
    """
    atomic_write_file(CERTS_FILE, json.dumps({"domains": sorted(domains)}, indent=2))


def add_cert_to_config(domain: str) -> None:
    """Add a domain to the certificate config.

    Args:
        domain: Domain name to add
    """
    with file_lock(f"{CERTS_FILE}.lock"):
        domains = load_certs_config()
        if domain not in domains:
            domains.append(domain)
            save_certs_config(domains)


def remove_cert_from_config(domain: str) -> None:
    """Remove a domain from the certificate config.

    Args:
        domain: Domain name to remove
    """
    with file_lock(f"{CERTS_FILE}.lock"):
        domains = load_certs_config()
        if domain in domains:
            domains.remove(domain)
            save_certs_config(domains)