"""File I/O operations for HAProxy MCP Server.""" import fcntl import json import os import tempfile from typing import Any, Optional from .config import ( MAP_FILE, WILDCARDS_MAP_FILE, SERVERS_FILE, CERTS_FILE, logger, ) from .validation import domain_to_backend 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 entries to separate map files for 2-stage matching. Uses 2-stage matching for performance: - domains.map: Exact domain matches (used with map_str, O(log n)) - wildcards.map: Wildcard entries (used with map_dom, O(n)) Args: entries: List of (domain, backend) tuples to write Raises: IOError: If the file cannot be written """ # 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 pool). Args: backend: Backend name to check Returns: True if this is a legacy backend, False if it's a pool """ 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 """ lock_path = f"{SERVERS_FILE}.lock" with open(lock_path, 'w') as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) try: 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) finally: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) 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 """ lock_path = f"{SERVERS_FILE}.lock" with open(lock_path, 'w') as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) try: 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) finally: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) def remove_domain_from_config(domain: str) -> None: """Remove domain from persistent config with file locking. Args: domain: Domain name to remove """ lock_path = f"{SERVERS_FILE}.lock" with open(lock_path, 'w') as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) try: config = load_servers_config() if domain in config: del config[domain] save_servers_config(config) finally: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) # 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 """ lock_path = f"{CERTS_FILE}.lock" with open(lock_path, 'w') as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) try: domains = load_certs_config() if domain not in domains: domains.append(domain) save_certs_config(domains) finally: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) def remove_cert_from_config(domain: str) -> None: """Remove a domain from the certificate config. Args: domain: Domain name to remove """ lock_path = f"{CERTS_FILE}.lock" with open(lock_path, 'w') as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) try: domains = load_certs_config() if domain in domains: domains.remove(domain) save_certs_config(domains) finally: fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)