MySQL Navigator MCP

import re import logging from typing import Dict, Any, Optional from pathlib import Path import ssl from functools import wraps import time # Import the Pydantic model from mcp_db.db import DatabaseCredentials logger = logging.getLogger(__name__) class SecurityManager: _instance = None def __new__(cls): if cls._instance is None: cls._instance = super(SecurityManager, cls).__new__(cls) cls._instance._initialize() return cls._instance def _initialize(self): """Initialize security settings""" self.rate_limit = { 'max_queries_per_minute': 100, 'current_minute': None, 'current_count': 0 } # Keep blocked patterns for query sanitization if needed elsewhere self.blocked_patterns = [ r'(?i)(?:delete|drop|truncate)\s+(?:table|database)', r'(?i)(?:alter|create)\s+(?:table|database|user)', r'(?i)execute\s+(?:procedure|function)', r'(?i)grant\s+|revoke\s+', r'(?i)system_user\(\)', r'(?i)super\s+privilege', ] def get_ssl_context(self, ssl_ca: Optional[str] = None) -> Optional[Dict[str, Any]]: """Create SSL context for secure connection""" if not ssl_ca: logger.debug("No SSL CA provided, SSL context not created.") return None try: ssl_ca_path = Path(ssl_ca) if not ssl_ca_path.is_file(): # Log as an error, but let the connection attempt handle the failure logger.error(f"SSL CA file not found at path: {ssl_ca}") raise FileNotFoundError(f"SSL CA file not found: {ssl_ca}") # Configuration options for mysql.connector's SSL arguments ssl_options = { 'ssl_ca': str(ssl_ca_path), # Add other options like ssl_cert, ssl_key, ssl_verify_cert if needed # 'ssl_verify_cert': True, # Often default, but can be explicit } logger.info(f"SSL context options prepared using CA: {ssl_ca}") return ssl_options except Exception as e: logger.error(f"Failed to prepare SSL context options: {e}", exc_info=True) # Propagate the error to be handled during connection attempt raise def sanitize_query(self, query: str) -> str: """Sanitize SQL query (basic checks for dangerous patterns).""" # Ensure query is a string if not isinstance(query, str): raise ValueError("Query must be a string.") # Remove comments (basic removal) query = re.sub(r'/\*.*?\*/', '', query, flags=re.DOTALL) # Handle multi-line comments query = re.sub(r'--.*$', '', query, flags=re.MULTILINE) query = re.sub(r'#.*$', '', query, flags=re.MULTILINE) # Handle hash comments # Check for dangerous patterns for pattern in self.blocked_patterns: if re.search(pattern, query, re.IGNORECASE): logger.warning(f"Query blocked due to forbidden pattern: {pattern}") raise ValueError(f"Query contains forbidden pattern: {pattern}") # Trim whitespace sanitized_query = query.strip() if not sanitized_query: raise ValueError("Query cannot be empty after sanitization.") logger.debug("Query passed basic sanitization checks.") return sanitized_query def check_rate_limit(self) -> bool: """Check if current request exceeds rate limit""" current_time = time.time() current_minute = int(current_time / 60) if self.rate_limit['current_minute'] != current_minute: # Reset count for the new minute self.rate_limit['current_minute'] = current_minute self.rate_limit['current_count'] = 0 logger.debug(f"Rate limit minute reset to {current_minute}") self.rate_limit['current_count'] += 1 logger.debug(f"Rate limit count for minute {current_minute}: {self.rate_limit['current_count']}/{self.rate_limit['max_queries_per_minute']}") if self.rate_limit['current_count'] > self.rate_limit['max_queries_per_minute']: logger.warning(f"Rate limit exceeded ({self.rate_limit['current_count']}/{self.rate_limit['max_queries_per_minute']})") return False # Rate limit exceeded return True # Within rate limit def sanitize_connection_params(self, params: DatabaseCredentials) -> None: """ Validate database connection parameters directly on the Pydantic model. Raises ValueError for invalid parameters. Note: This method modifies the validation logic to work with the model's attributes. It no longer returns a dict, as it validates the input model directly. """ logger.debug(f"Sanitizing connection parameters for host: {params.host}, db: {params.database}") # Validate host (allow hostname or IP address) # More permissive regex: allows domain names, localhost, IPv4, IPv6 host_pattern = r"^(?:(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,6}|localhost|(?:\d{1,3}\.){3}\d{1,3}|\[?[a-fA-F0-9:]+\]?)$" if not re.match(host_pattern, params.host): logger.error(f"Invalid host format detected: {params.host}") raise ValueError(f"Invalid host format: {params.host}") # Validate port if not isinstance(params.port, int) or not (1 <= params.port <= 65535): logger.error(f"Invalid port number detected: {params.port}") raise ValueError(f"Invalid port number: {params.port}. Must be between 1 and 65535.") # Validate database name (allow alphanumeric, underscore, hyphen) db_name_pattern = r'^[a-zA-Z0-9_-]+$' if not re.match(db_name_pattern, params.database): logger.error(f"Invalid database name format detected: {params.database}") raise ValueError(f"Invalid database name format: {params.database}. Use alphanumeric, underscore, or hyphen.") # Validate username (similar to database name) username_pattern = r'^[a-zA-Z0-9_@.-]+$' # Allow more chars often used in usernames if not re.match(username_pattern, params.username): logger.error(f"Invalid username format detected: {params.username}") raise ValueError(f"Invalid username format: {params.username}") # Validate password (check if it's set - Pydantic handles SecretStr) if not params.password or not params.password.get_secret_value(): logger.error("Database password is required but not provided.") raise ValueError("Database password is required.") # Validate SSL CA path if provided if params.ssl_ca: ssl_ca_path = Path(params.ssl_ca) if not ssl_ca_path.is_file(): # Log as warning, connection will fail later if file is truly needed and missing logger.warning(f"Provided SSL CA path does not point to a file: {params.ssl_ca}") # Optionally raise ValueError here if SSL is mandatory and file must exist # raise ValueError(f"SSL CA file not found at path: {params.ssl_ca}") # Validate retries and delay if not isinstance(params.max_retries, int) or params.max_retries < 0: logger.error(f"Invalid max_retries value: {params.max_retries}") raise ValueError("max_retries must be a non-negative integer.") if not isinstance(params.retry_delay, (int, float)) or params.retry_delay < 0: logger.error(f"Invalid retry_delay value: {params.retry_delay}") raise ValueError("retry_delay must be a non-negative number.") logger.debug("Connection parameters passed validation.") # No need to return anything as we are validating the passed object. def enforce_security(func): """Decorator to enforce security measures like rate limiting and query sanitization.""" @wraps(func) def wrapper(*args, **kwargs): # Get the SecurityManager instance # Assumes SecurityManager is initialized elsewhere (e.g., in main.py) # If not, might need a way to access the instance created in main.py # For simplicity here, we re-create it, but a shared instance is better. security_manager = SecurityManager() # Or access a shared instance # --- Check Rate Limit --- if not security_manager.check_rate_limit(): # Raise an exception that the handle_error decorator can catch raise Exception("Rate limit exceeded. Please try again later.") # Or a custom RateLimitError # --- Sanitize Query (if applicable) --- # Check if 'query' or 'query_params' exists in kwargs and needs sanitization # This part depends on which tool is being decorated and its arguments. # Example: If decorating query_database which takes query_params dict if func.__name__ == 'query_database' and 'query_params' in kwargs: # This example assumes the raw SQL isn't passed directly. # If a raw SQL string *was* passed (which is discouraged), sanitize it: # if 'raw_sql_query' in kwargs: # kwargs['raw_sql_query'] = security_manager.sanitize_query(kwargs['raw_sql_query']) # Since we build parameterized queries, direct SQL sanitization might not be needed here, # but validation of inputs (table names, field names) within query_database is crucial. logger.debug(f"Security checks passed for {func.__name__}.") # Placeholder log pass # Parameterized queries handle SQL injection risk # Execute the original function return func(*args, **kwargs) # Preserve original function signature for MCP introspection wrapper.__name__ = func.__name__ wrapper.__doc__ = func.__doc__ wrapper.__annotations__ = func.__annotations__ return wrapper