mcp-security-framework 0.1.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

mcp_security_framework/middleware/rate_limit_middleware.py

@@ -30,7 +30,7 @@ from .security_middleware import SecurityMiddleware, SecurityMiddlewareError
 
 class RateLimitMiddlewareError(SecurityMiddlewareError):
     """Raised when rate limit middleware encounters an error."""
-    
+
     def __init__(self, message: str, error_code: int = -32035):
         self.message = message
         self.error_code = error_code
@@ -40,161 +40,167 @@ class RateLimitMiddlewareError(SecurityMiddlewareError):
 class RateLimitMiddleware(SecurityMiddleware):
     """
     Rate Limiting-Only Middleware Class
-    
+
     This class provides rate limiting-only middleware that focuses
     solely on request rate limiting without performing authentication
     or authorization checks. It's useful for scenarios where rate
     limiting is handled separately from other security concerns.
-    
+
     The RateLimitMiddleware implements:
     - Rate limiting-only request processing
     - Multiple rate limiting strategies
     - Rate limit result caching
     - Rate limit event logging
     - Framework-agnostic design
-    
+
     Key Responsibilities:
     - Process requests through rate limiting pipeline only
     - Handle multiple rate limiting strategies
     - Cache rate limit results for performance
     - Log rate limit events and violations
     - Provide rate limit status to downstream components
-    
+
     Attributes:
         Inherits all attributes from SecurityMiddleware
         _rate_limit_cache (Dict): Cache for rate limit results
         _rate_limit_strategies (Dict): Available rate limiting strategies
-        
+
     Example:
         >>> from mcp_security_framework.middleware import RateLimitMiddleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> rate_limit_middleware = RateLimitMiddleware(security_manager)
         >>> app.add_middleware(rate_limit_middleware)
-        
+
     Note:
         This middleware only handles rate limiting. Authentication and
         authorization should be handled separately by other middleware
         or application logic.
     """
-    
+
     def __init__(self, security_manager):
         """
         Initialize Rate Limiting-Only Middleware.
-        
+
         Args:
             security_manager: Security manager instance containing
                 all security components and configuration.
-                
+
         Raises:
             RateLimitMiddlewareError: If initialization fails
         """
         super().__init__(security_manager)
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
-        
+
         # Initialize rate limiting strategies
         self._rate_limit_strategies = {
             "ip": self._rate_limit_by_ip,
             "user": self._rate_limit_by_user,
             "global": self._rate_limit_global,
             "path": self._rate_limit_by_path,
-            "method": self._rate_limit_by_method
+            "method": self._rate_limit_by_method,
         }
-        
+
         self.logger.info("Rate limit middleware initialized")
-    
+
     @abstractmethod
     def __call__(self, request: Any, call_next: Any) -> Any:
         """
         Process request through rate limiting middleware.
-        
+
         This method implements the rate limiting-only processing
         pipeline, focusing solely on request rate limiting.
-        
+
         Args:
             request: Framework-specific request object
             call_next: Framework-specific call_next function
-            
+
         Returns:
             Framework-specific response object
-            
+
         Raises:
             RateLimitMiddlewareError: If rate limiting processing fails
         """
         pass
-    
+
     def _check_rate_limit_only(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Perform rate limiting-only processing.
-        
+
         This method handles rate limiting without authentication
         or authorization checks.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
-            
+
         Raises:
             RateLimitMiddlewareError: If rate limiting process fails
         """
         try:
             if not self.config.rate_limit.enabled:
-                return True, {"strategy": "disabled", "reason": "Rate limiting disabled"}
-            
+                return True, {
+                    "strategy": "disabled",
+                    "reason": "Rate limiting disabled",
+                }
+
             # Get rate limiting strategy from config
-            strategy = self.config.rate_limit.strategy if hasattr(self.config.rate_limit, 'strategy') else "ip"
-            
+            strategy = (
+                self.config.rate_limit.strategy
+                if hasattr(self.config.rate_limit, "strategy")
+                else "ip"
+            )
+
             if strategy not in self._rate_limit_strategies:
-                self.logger.warning(f"Unknown rate limiting strategy: {strategy}, using 'ip'")
+                self.logger.warning(
+                    f"Unknown rate limiting strategy: {strategy}, using 'ip'"
+                )
                 strategy = "ip"
-            
+
             # Apply rate limiting strategy
             is_allowed, rate_limit_info = self._rate_limit_strategies[strategy](request)
-            
+
             # Log rate limit event
             self._log_rate_limit_event(request, is_allowed, rate_limit_info)
-            
+
             return is_allowed, rate_limit_info
-            
+
         except Exception as e:
             self.logger.error(
-                "Rate limiting process failed",
-                extra={"error": str(e)},
-                exc_info=True
+                "Rate limiting process failed", extra={"error": str(e)}, exc_info=True
             )
             raise RateLimitMiddlewareError(
-                f"Rate limiting process failed: {str(e)}",
-                error_code=-32036
+                f"Rate limiting process failed: {str(e)}", error_code=-32036
             )
-    
+
     def _rate_limit_by_ip(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Rate limit by IP address.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
         """
         identifier = self._get_rate_limit_identifier(request)
         is_allowed = self.security_manager.rate_limiter.check_rate_limit(identifier)
-        
+
         return is_allowed, {
             "strategy": "ip",
             "identifier": identifier,
-            "is_allowed": is_allowed
+            "is_allowed": is_allowed,
         }
-    
+
     def _rate_limit_by_user(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Rate limit by user (requires authentication).
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
         """
@@ -203,114 +209,117 @@ class RateLimitMiddleware(SecurityMiddleware):
         if not user_id:
             # Fall back to IP-based rate limiting
             return self._rate_limit_by_ip(request)
-        
-        is_allowed = self.security_manager.rate_limiter.check_rate_limit(f"user:{user_id}")
-        
+
+        is_allowed = self.security_manager.rate_limiter.check_rate_limit(
+            f"user:{user_id}"
+        )
+
         return is_allowed, {
             "strategy": "user",
             "identifier": user_id,
-            "is_allowed": is_allowed
+            "is_allowed": is_allowed,
         }
-    
+
     def _rate_limit_global(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Global rate limiting (all requests).
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
         """
         is_allowed = self.security_manager.rate_limiter.check_rate_limit("global")
-        
+
         return is_allowed, {
             "strategy": "global",
             "identifier": "global",
-            "is_allowed": is_allowed
+            "is_allowed": is_allowed,
         }
-    
+
     def _rate_limit_by_path(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Rate limit by request path.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
         """
         path = self._get_request_path(request)
         identifier = f"path:{path}"
         is_allowed = self.security_manager.rate_limiter.check_rate_limit(identifier)
-        
+
         return is_allowed, {
             "strategy": "path",
             "identifier": identifier,
             "path": path,
-            "is_allowed": is_allowed
+            "is_allowed": is_allowed,
         }
-    
+
     def _rate_limit_by_method(self, request: Any) -> Tuple[bool, Dict[str, Any]]:
         """
         Rate limit by HTTP method.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Tuple[bool, Dict[str, Any]]: (is_allowed, rate_limit_info)
         """
         method = self._get_request_method(request)
         identifier = f"method:{method}"
         is_allowed = self.security_manager.rate_limiter.check_rate_limit(identifier)
-        
+
         return is_allowed, {
             "strategy": "method",
             "identifier": identifier,
             "method": method,
-            "is_allowed": is_allowed
+            "is_allowed": is_allowed,
         }
-    
+
     def _get_user_identifier(self, request: Any) -> Optional[str]:
         """
         Get user identifier from request.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             Optional[str]: User identifier if available, None otherwise
         """
         # This should be implemented by framework-specific subclasses
         # For now, return None to indicate no user identifier available
         return None
-    
+
     def _get_request_method(self, request: Any) -> str:
         """
         Get HTTP method from request.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             str: HTTP method
         """
         # This should be implemented by framework-specific subclasses
         return "GET"
-    
-    def _log_rate_limit_event(self, request: Any, is_allowed: bool, 
-                             rate_limit_info: Dict[str, Any]) -> None:
+
+    def _log_rate_limit_event(
+        self, request: Any, is_allowed: bool, rate_limit_info: Dict[str, Any]
+    ) -> None:
         """
         Log rate limit event.
-        
+
         Args:
             request: Framework-specific request object
             is_allowed (bool): Whether request is allowed
             rate_limit_info (Dict[str, Any]): Rate limit information
         """
         log_level = logging.WARNING if not is_allowed else logging.DEBUG
-        
+
         self.logger.log(
             log_level,
             f"Rate limit event: {'allowed' if is_allowed else 'blocked'}",
@@ -320,68 +329,64 @@ class RateLimitMiddleware(SecurityMiddleware):
                 "ip_address": self._get_rate_limit_identifier(request),
                 "path": self._get_request_path(request),
                 "method": self._get_request_method(request),
-                **rate_limit_info
-            }
+                **rate_limit_info,
+            },
         )
-    
+
     def get_rate_limit_status(self, identifier: str) -> Dict[str, Any]:
         """
         Get current rate limit status for an identifier.
-        
+
         Args:
             identifier (str): Rate limit identifier
-            
+
         Returns:
             Dict[str, Any]: Rate limit status information
         """
         try:
-            status = self.security_manager.rate_limiter.get_rate_limit_status(identifier)
+            status = self.security_manager.rate_limiter.get_rate_limit_status(
+                identifier
+            )
             return {
                 "identifier": identifier,
                 "current_count": status.current_count,
                 "limit": status.limit,
                 "window_seconds": status.window_seconds,
                 "remaining": status.remaining,
-                "reset_time": status.reset_time
+                "reset_time": status.reset_time,
             }
         except Exception as e:
             self.logger.error(
                 "Failed to get rate limit status",
-                extra={"identifier": identifier, "error": str(e)}
+                extra={"identifier": identifier, "error": str(e)},
             )
-            return {
-                "identifier": identifier,
-                "error": str(e)
-            }
-    
+            return {"identifier": identifier, "error": str(e)}
+
     def reset_rate_limit(self, identifier: str) -> bool:
         """
         Reset rate limit for an identifier.
-        
+
         Args:
             identifier (str): Rate limit identifier
-            
+
         Returns:
             bool: True if reset successful, False otherwise
         """
         try:
             self.security_manager.rate_limiter.reset_rate_limit(identifier)
-            self.logger.info(
-                "Rate limit reset",
-                extra={"identifier": identifier}
-            )
+            self.logger.info("Rate limit reset", extra={"identifier": identifier})
             return True
         except Exception as e:
             self.logger.error(
                 "Failed to reset rate limit",
-                extra={"identifier": identifier, "error": str(e)}
+                extra={"identifier": identifier, "error": str(e)},
             )
             return False
-    
+
     def get_rate_limit_statistics(self) -> Dict[str, Any]:
         """
         Get rate limiting statistics.
-        
+
         Returns:
             Dict[str, Any]: Rate limiting statistics
         """
@@ -393,11 +398,10 @@ class RateLimitMiddleware(SecurityMiddleware):
                 "blocked_requests": 0,
                 "active_identifiers": 0,
                 "cache_hits": 0,
-                "cache_misses": 0
+                "cache_misses": 0,
             }
         except Exception as e:
             self.logger.error(
-                "Failed to get rate limit statistics",
-                extra={"error": str(e)}
+                "Failed to get rate limit statistics", extra={"error": str(e)}
             )
             return {"error": str(e)}