rootly-mcp-server 2.0.14__py3-none-any.whl → 2.1.0__py3-none-any.whl

rootly_mcp_server/security.py

@@ -0,0 +1,404 @@
+"""
+Security utilities for the Rootly MCP Server.
+
+This module provides security-related functionality including:
+- Secure token handling
+- HTTPS enforcement
+- Input sanitization
+- Rate limiting
+- Security validation
+"""
+
+import os
+import re
+import time
+from collections import defaultdict
+from functools import wraps
+from threading import Lock
+from typing import Any
+from urllib.parse import urlparse
+
+from .exceptions import (
+    RootlyConfigurationError,
+    RootlyRateLimitError,
+    RootlyValidationError,
+)
+
+# Token validation pattern (Bearer tokens typically start with a prefix)
+TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9_-]{20,}$")
+
+# URL validation patterns
+HTTPS_PATTERN = re.compile(r"^https://")
+VALID_DOMAIN_PATTERN = re.compile(
+    r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,}$"
+)
+
+# SQL injection patterns to block
+SQL_INJECTION_PATTERNS = [
+    r"(\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|EXEC|EXECUTE|UNION|SCRIPT)\b)",
+    r"(--|;|\/\*|\*\/|xp_|sp_)",
+    r"(\bOR\b.*=.*)",
+    r"(\bAND\b.*=.*)",
+]
+
+# XSS patterns to block
+XSS_PATTERNS = [
+    r"<script[^>]*>.*?</script>",
+    r"javascript:",
+    r"onerror\s*=",
+    r"onload\s*=",
+    r"<iframe[^>]*>",
+]
+
+
+class RateLimiter:
+    """
+    Token bucket rate limiter for API requests.
+
+    Implements a sliding window rate limiter to prevent API abuse.
+    """
+
+    def __init__(self, max_requests: int = 100, time_window: int = 60):
+        """
+        Initialize the rate limiter.
+
+        Args:
+            max_requests: Maximum number of requests allowed in the time window
+            time_window: Time window in seconds
+        """
+        self.max_requests = max_requests
+        self.time_window = time_window
+        self._requests = defaultdict(list)
+        self._lock = Lock()
+
+    def is_allowed(self, identifier: str) -> tuple[bool, int | None]:
+        """
+        Check if a request is allowed for the given identifier.
+
+        Args:
+            identifier: Unique identifier for the client (e.g., IP address, user ID)
+
+        Returns:
+            Tuple of (is_allowed, retry_after_seconds)
+        """
+        with self._lock:
+            current_time = time.time()
+            window_start = current_time - self.time_window
+
+            # Clean up old requests
+            self._requests[identifier] = [
+                req_time for req_time in self._requests[identifier] if req_time > window_start
+            ]
+
+            # Check if limit is exceeded
+            if len(self._requests[identifier]) >= self.max_requests:
+                # Calculate retry_after based on oldest request
+                oldest_request = min(self._requests[identifier])
+                retry_after = int(oldest_request + self.time_window - current_time) + 1
+                return False, retry_after
+
+            # Allow the request and record it
+            self._requests[identifier].append(current_time)
+            return True, None
+
+    def reset(self, identifier: str) -> None:
+        """Reset the rate limit for a specific identifier."""
+        with self._lock:
+            self._requests.pop(identifier, None)
+
+
+# Global rate limiter instance
+_rate_limiter = RateLimiter(max_requests=100, time_window=60)
+
+
+def get_rate_limiter() -> RateLimiter:
+    """Get the global rate limiter instance."""
+    return _rate_limiter
+
+
+def rate_limit(identifier_func=None):
+    """
+    Decorator to apply rate limiting to a function.
+
+    Args:
+        identifier_func: Optional function to extract identifier from function args.
+                        If None, uses "default" as identifier.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            identifier = identifier_func(*args, **kwargs) if identifier_func else "default"
+            allowed, retry_after = _rate_limiter.is_allowed(identifier)
+
+            if not allowed:
+                raise RootlyRateLimitError(
+                    f"Rate limit exceeded. Try again in {retry_after} seconds.",
+                    retry_after=retry_after,
+                )
+
+            return await func(*args, **kwargs)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            identifier = identifier_func(*args, **kwargs) if identifier_func else "default"
+            allowed, retry_after = _rate_limiter.is_allowed(identifier)
+
+            if not allowed:
+                raise RootlyRateLimitError(
+                    f"Rate limit exceeded. Try again in {retry_after} seconds.",
+                    retry_after=retry_after,
+                )
+
+            return func(*args, **kwargs)
+
+        # Return appropriate wrapper based on function type
+        import asyncio
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        return sync_wrapper
+
+    return decorator
+
+
+def validate_api_token(token: str | None) -> str:
+    """
+    Validate that an API token is properly formatted and not empty.
+
+    Args:
+        token: The API token to validate
+
+    Returns:
+        The validated token
+
+    Raises:
+        RootlyConfigurationError: If token is missing or invalid
+    """
+    if not token:
+        raise RootlyConfigurationError(
+            "API token is required but not provided. Set the ROOTLY_API_TOKEN environment variable."
+        )
+
+    token = token.strip()
+
+    if len(token) < 20:
+        raise RootlyConfigurationError(
+            "API token appears to be invalid (too short). Please check your ROOTLY_API_TOKEN value."
+        )
+
+    # Don't log the actual token value for security
+    return token
+
+
+def get_api_token_from_env() -> str:
+    """
+    Get and validate the API token from environment variables.
+
+    Returns:
+        The validated API token
+
+    Raises:
+        RootlyConfigurationError: If token is missing or invalid
+    """
+    token = os.getenv("ROOTLY_API_TOKEN")
+    return validate_api_token(token)
+
+
+def enforce_https(url: str) -> str:
+    """
+    Ensure that a URL uses HTTPS.
+
+    Args:
+        url: The URL to validate
+
+    Returns:
+        The validated HTTPS URL
+
+    Raises:
+        RootlyValidationError: If URL doesn't use HTTPS
+    """
+    if not url:
+        raise RootlyValidationError("URL cannot be empty")
+
+    parsed = urlparse(url)
+
+    if not parsed.scheme:
+        # Assume HTTPS if no scheme provided
+        return f"https://{url}"
+
+    if parsed.scheme != "https":
+        raise RootlyValidationError(
+            f"Only HTTPS URLs are allowed for security reasons. Got: {parsed.scheme}://"
+        )
+
+    return url
+
+
+def validate_url(url: str, allowed_domains: list[str] | None = None) -> str:
+    """
+    Validate a URL for security.
+
+    Args:
+        url: The URL to validate
+        allowed_domains: Optional list of allowed domains
+
+    Returns:
+        The validated URL
+
+    Raises:
+        RootlyValidationError: If URL is invalid or not allowed
+    """
+    if not url:
+        raise RootlyValidationError("URL cannot be empty")
+
+    # Enforce HTTPS
+    url = enforce_https(url)
+
+    parsed = urlparse(url)
+
+    # Validate domain
+    if not parsed.netloc:
+        raise RootlyValidationError(f"Invalid URL: missing domain in {url}")
+
+    # Check against allowed domains if provided
+    if allowed_domains:
+        domain_allowed = any(
+            parsed.netloc == domain or parsed.netloc.endswith(f".{domain}")
+            for domain in allowed_domains
+        )
+        if not domain_allowed:
+            raise RootlyValidationError(f"Domain {parsed.netloc} is not in the allowed list")
+
+    return url
+
+
+def sanitize_input(value: Any, max_length: int = 10000) -> Any:
+    """
+    Sanitize user input to prevent injection attacks.
+
+    Args:
+        value: The value to sanitize
+        max_length: Maximum allowed length for strings
+
+    Returns:
+        The sanitized value
+
+    Raises:
+        RootlyValidationError: If input contains malicious patterns
+    """
+    if value is None:
+        return None
+
+    if isinstance(value, bool):
+        return value
+
+    if isinstance(value, int | float):
+        return value
+
+    if isinstance(value, str):
+        # Check length
+        if len(value) > max_length:
+            raise RootlyValidationError(
+                f"Input too long: {len(value)} characters (max: {max_length})"
+            )
+
+        # Check for SQL injection patterns
+        for pattern in SQL_INJECTION_PATTERNS:
+            if re.search(pattern, value, re.IGNORECASE):
+                raise RootlyValidationError("Input contains potentially malicious SQL patterns")
+
+        # Check for XSS patterns
+        for pattern in XSS_PATTERNS:
+            if re.search(pattern, value, re.IGNORECASE):
+                raise RootlyValidationError("Input contains potentially malicious XSS patterns")
+
+        return value
+
+    if isinstance(value, dict):
+        return {k: sanitize_input(v, max_length) for k, v in value.items()}
+
+    if isinstance(value, list | tuple):
+        return type(value)(sanitize_input(item, max_length) for item in value)
+
+    # For other types, convert to string and sanitize
+    return sanitize_input(str(value), max_length)
+
+
+def sanitize_error_message(error_message: str, max_length: int = 500) -> str:
+    """
+    Sanitize error messages to prevent information leakage.
+
+    Removes file paths, stack traces, and other sensitive information.
+
+    Args:
+        error_message: The error message to sanitize
+        max_length: Maximum length for the sanitized message
+
+    Returns:
+        The sanitized error message
+    """
+    if not error_message:
+        return "An error occurred"
+
+    # Remove absolute file paths
+    error_message = re.sub(r"/[\w/.-]+\.py", "[file]", error_message)
+    error_message = re.sub(r"C:\\[\w\\.-]+\.py", "[file]", error_message)
+
+    # Remove line numbers
+    error_message = re.sub(r", line \d+", "", error_message)
+
+    # Remove "Traceback" and everything after it
+    if "Traceback" in error_message:
+        error_message = error_message.split("Traceback")[0].strip()
+
+    # Remove stack trace markers
+    error_message = re.sub(r"File \"[^\"]+\"", "", error_message)
+
+    # Truncate if too long
+    if len(error_message) > max_length:
+        error_message = error_message[:max_length] + "..."
+
+    return error_message.strip() or "An error occurred"
+
+
+def mask_sensitive_data(
+    data: dict[str, Any], sensitive_keys: list[str] | None = None
+) -> dict[str, Any]:
+    """
+    Mask sensitive data in a dictionary for logging.
+
+    Args:
+        data: The data dictionary to mask
+        sensitive_keys: List of key patterns to mask (case-insensitive)
+
+    Returns:
+        Dictionary with sensitive values masked
+    """
+    if sensitive_keys is None:
+        sensitive_keys = ["token", "password", "secret", "api_key", "auth"]
+
+    def should_mask(key: str) -> bool:
+        key_lower = key.lower()
+        return any(sensitive in key_lower for sensitive in sensitive_keys)
+
+    def mask_value(value: Any) -> Any:
+        if isinstance(value, str) and len(value) > 0:
+            return "***REDACTED***"
+        return value
+
+    masked = {}
+    for key, value in data.items():
+        if should_mask(key):
+            masked[key] = mask_value(value)
+        elif isinstance(value, dict):
+            masked[key] = mask_sensitive_data(value, sensitive_keys)
+        elif isinstance(value, list):
+            masked[key] = [
+                mask_sensitive_data(item, sensitive_keys) if isinstance(item, dict) else item
+                for item in value
+            ]
+        else:
+            masked[key] = value
+
+    return masked