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

mcp_security_framework/constants.py

@@ -62,16 +62,17 @@ HTTP_NOT_FOUND = 404
 HTTP_TOO_MANY_REQUESTS = 429
 HTTP_INTERNAL_SERVER_ERROR = 500
 
+
 # Error Codes
 class ErrorCodes:
     """Error codes for the MCP Security Framework."""
-    
+
     # SSL/TLS Errors (-32001 to -32009)
     SSL_CONFIGURATION_ERROR = -32001
     CERTIFICATE_VALIDATION_ERROR = -32002
     SSL_CONTEXT_CREATION_ERROR = -32003
     SSL_HANDSHAKE_ERROR = -32004
-    
+
     # Authentication Errors (-32010 to -32019)
     AUTHENTICATION_ERROR = -32010
     AUTHENTICATION_CONFIGURATION_ERROR = -32011
@@ -80,46 +81,47 @@ class ErrorCodes:
     CERTIFICATE_AUTH_ERROR = -32014
     BASIC_AUTH_ERROR = -32015
     AUTH_METHOD_NOT_SUPPORTED = -32016
-    
+
     # Authorization Errors (-32020 to -32029)
     PERMISSION_DENIED_ERROR = -32020
     INSUFFICIENT_PERMISSIONS = -32021
     ROLE_NOT_FOUND = -32022
     PERMISSION_NOT_FOUND = -32023
-    
+
     # Rate Limiting Errors (-32030 to -32039)
     RATE_LIMIT_EXCEEDED_ERROR = -32030
     RATE_LIMIT_CONFIGURATION_ERROR = -32031
     RATE_LIMIT_STORAGE_ERROR = -32032
-    
+
     # Middleware Errors (-32040 to -32049)
     SECURITY_MIDDLEWARE_ERROR = -32040
     AUTH_MIDDLEWARE_ERROR = -32041
     MTLS_MIDDLEWARE_ERROR = -32042
     RATE_LIMIT_MIDDLEWARE_ERROR = -32043
-    
+
     # Certificate Management Errors (-32050 to -32059)
     CERTIFICATE_GENERATION_ERROR = -32050
     CERTIFICATE_REVOCATION_ERROR = -32051
     CERTIFICATE_STORAGE_ERROR = -32052
     CA_CONFIGURATION_ERROR = -32053
-    
+
     # Crypto Errors (-32060 to -32069)
     CRYPTO_ERROR = -32060
     KEY_GENERATION_ERROR = -32061
     HASHING_ERROR = -32062
     ENCRYPTION_ERROR = -32063
-    
+
     # Configuration Errors (-32070 to -32079)
     CONFIGURATION_ERROR = -32070
     VALIDATION_ERROR = -32071
     SERIALIZATION_ERROR = -32072
-    
+
     # General Errors (-32080 to -32099)
     GENERAL_ERROR = -32080
     NOT_IMPLEMENTED_ERROR = -32081
     UNSUPPORTED_OPERATION_ERROR = -32082
 
+
 # Security Headers
 DEFAULT_SECURITY_HEADERS = {
     "X-Content-Type-Options": "nosniff",
@@ -127,7 +129,7 @@ DEFAULT_SECURITY_HEADERS = {
     "X-XSS-Protection": "1; mode=block",
     "Strict-Transport-Security": "max-age=31536000; includeSubDomains",
     "Content-Security-Policy": "default-src 'self'",
-    "Referrer-Policy": "strict-origin-when-cross-origin"
+    "Referrer-Policy": "strict-origin-when-cross-origin",
 }
 
 # Authentication Methods
@@ -136,29 +138,21 @@ AUTH_METHODS = {
     "JWT": "jwt",
     "CERTIFICATE": "certificate",
     "BASIC": "basic",
-    "OAUTH2": "oauth2"
+    "OAUTH2": "oauth2",
 }
 
 # Storage Backends
-STORAGE_BACKENDS = {
-    "MEMORY": "memory",
-    "REDIS": "redis",
-    "DATABASE": "database"
-}
+STORAGE_BACKENDS = {"MEMORY": "memory", "REDIS": "redis", "DATABASE": "database"}
 
 # Hash Algorithms
-HASH_ALGORITHMS = {
-    "SHA256": "sha256",
-    "SHA512": "sha512",
-    "MD5": "md5"
-}
+HASH_ALGORITHMS = {"SHA256": "sha256", "SHA512": "sha512", "MD5": "md5"}
 
 # TLS Versions
 TLS_VERSIONS = {
     "TLSv1.0": "TLSv1.0",
     "TLSv1.1": "TLSv1.1",
     "TLSv1.2": "TLSv1.2",
-    "TLSv1.3": "TLSv1.3"
+    "TLSv1.3": "TLSv1.3",
 }
 
 # Certificate Revocation Reasons
@@ -169,7 +163,7 @@ CERTIFICATE_REVOCATION_REASONS = {
     "AFFILIATION_CHANGED": "affiliation_changed",
     "SUPERSEDED": "superseded",
     "CESSATION_OF_OPERATION": "cessation_of_operation",
-    "CERTIFICATE_HOLD": "certificate_hold"
+    "CERTIFICATE_HOLD": "certificate_hold",
 }
 
 # Log Levels
@@ -178,7 +172,7 @@ LOG_LEVELS = {
     "INFO": "INFO",
     "WARNING": "WARNING",
     "ERROR": "ERROR",
-    "CRITICAL": "CRITICAL"
+    "CRITICAL": "CRITICAL",
 }
 
 # Time Constants (in seconds)
@@ -188,7 +182,7 @@ TIME_CONSTANTS = {
     "DAY": 86400,
     "WEEK": 604800,
     "MONTH": 2592000,  # 30 days
-    "YEAR": 31536000   # 365 days
+    "YEAR": 31536000,  # 365 days
 }
 
 # Cache Keys
@@ -196,7 +190,7 @@ CACHE_KEY_PREFIXES = {
     "AUTH": "auth",
     "RATE_LIMIT": "rate_limit",
     "PERMISSION": "permission",
-    "CERTIFICATE": "certificate"
+    "CERTIFICATE": "certificate",
 }
 
 # Environment Variables
@@ -205,5 +199,5 @@ ENV_VARS = {
     "DEFAULT_SERVER_HOST": "DEFAULT_SERVER_HOST",
     "DEFAULT_SERVER_PORT": "DEFAULT_SERVER_PORT",
     "LOG_LEVEL": "LOG_LEVEL",
-    "ENVIRONMENT": "ENVIRONMENT"
+    "ENVIRONMENT": "ENVIRONMENT",
 }

mcp_security_framework/core/auth_manager.py

@@ -26,7 +26,7 @@ License: MIT
 import logging
 import time
 from datetime import datetime, timedelta, timezone
-from typing import Dict, List, Optional, Union, Any
+from typing import Any, Dict, List, Optional, Union
 
 import jwt
 from cryptography import x509
@@ -47,6 +47,9 @@ from ..utils.crypto_utils import (
     validate_api_key_format,
     verify_password,
 )
+from ..utils.datetime_compat import (
+    get_not_valid_after_utc,
+)
 
 
 class AuthManager:
@@ -136,7 +139,9 @@ class AuthManager:
                     self._api_keys[username] = key
                     self._api_key_metadata[key] = value
                 else:
-                    self.logger.warning(f"Invalid API key format for key {key}: {value}")
+                    self.logger.warning(
+                        f"Invalid API key format for key {key}: {value}"
+                    )
         else:
             self._api_keys = {}
             self._api_key_metadata = {}
@@ -159,12 +164,12 @@ class AuthManager:
                 "jwt_expiry_hours": config.jwt_expiry_hours,
             },
         )
-    
+
     @property
     def is_auth_enabled(self) -> bool:
         """
         Check if authentication is enabled.
-        
+
         Returns:
             bool: True if authentication is enabled, False otherwise
         """
@@ -287,16 +292,24 @@ class AuthManager:
             user_permissions = set()
             if self.permission_manager:
                 try:
-                    permissions_result = self.permission_manager.get_effective_permissions(user_roles)
+                    permissions_result = (
+                        self.permission_manager.get_effective_permissions(user_roles)
+                    )
                     # Handle both set and mock objects
-                    if hasattr(permissions_result, '__iter__') and not isinstance(permissions_result, str):
+                    if hasattr(permissions_result, "__iter__") and not isinstance(
+                        permissions_result, str
+                    ):
                         user_permissions = set(permissions_result)
                     else:
                         user_permissions = set()
                 except Exception as e:
                     self.logger.warning(
                         "Failed to get user permissions",
-                        extra={"username": username, "roles": user_roles, "error": str(e)},
+                        extra={
+                            "username": username,
+                            "roles": user_roles,
+                            "error": str(e),
+                        },
                     )
 
             # Create successful authentication result
@@ -458,9 +471,13 @@ class AuthManager:
             user_permissions = set()
             if self.permission_manager:
                 try:
-                    permissions_result = self.permission_manager.get_effective_permissions(roles)
+                    permissions_result = (
+                        self.permission_manager.get_effective_permissions(roles)
+                    )
                     # Handle both set and mock objects
-                    if hasattr(permissions_result, '__iter__') and not isinstance(permissions_result, str):
+                    if hasattr(permissions_result, "__iter__") and not isinstance(
+                        permissions_result, str
+                    ):
                         user_permissions = set(permissions_result)
                     else:
                         user_permissions = set()
@@ -578,7 +595,7 @@ class AuthManager:
 
             # Check certificate expiration
             now = datetime.now(timezone.utc)
-            if cert.not_valid_after_utc.replace(tzinfo=timezone.utc) < now:
+            if get_not_valid_after_utc(cert) < now:
                 return AuthResult(
                     is_valid=False,
                     status=AuthStatus.EXPIRED,
@@ -636,7 +653,7 @@ class AuthManager:
                 roles=roles,
                 auth_method="certificate",
                 auth_timestamp=datetime.now(timezone.utc),
-                token_expiry=cert.not_valid_after_utc.replace(tzinfo=timezone.utc),
+                token_expiry=get_not_valid_after_utc(cert),
             )
 
             self.logger.info(
@@ -906,37 +923,39 @@ class AuthManager:
     def _get_user_roles(self, username: str) -> List[str]:
         """
         Get user roles from configuration.
-        
+
         This method retrieves user roles from the authentication configuration.
         It checks both the user_roles mapping and the permission manager.
-        
+
         Args:
             username (str): Username to get roles for
-            
+
         Returns:
             List[str]: List of user roles
         """
         # Check user_roles mapping first
         if self.config.user_roles and username in self.config.user_roles:
             return self.config.user_roles[username]
-        
+
         # Fallback to permission manager
         if self.permission_manager:
             return self.permission_manager.get_user_roles(username)
-        
+
         return []
 
-    def _validate_external_user(self, username: str, credentials: Dict[str, Any]) -> bool:
+    def _validate_external_user(
+        self, username: str, credentials: Dict[str, Any]
+    ) -> bool:
         """
         Validate user against external authentication system.
-        
+
         This method provides a hook for integrating with external
         authentication systems (LDAP, Active Directory, etc.).
-        
+
         Args:
             username (str): Username to validate
             credentials (Dict[str, Any]): User credentials
-            
+
         Returns:
             bool: True if user is valid in external system
         """
@@ -944,9 +963,9 @@ class AuthManager:
         # In a real implementation, this would connect to external auth systems
         self.logger.debug(
             "External user validation called",
-            extra={"username": username, "method": "external"}
+            extra={"username": username, "method": "external"},
         )
-        
+
         # For now, return False to indicate external validation not implemented
         return False