mcp-security-framework 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

mcp_security_framework/__init__.py

@@ -71,7 +71,7 @@ from mcp_security_framework.schemas.responses import (
 )
 
 # Version information
-__version__ = "0.1.0"
+__version__ = "1.2.1"
 __author__ = "Vasiliy Zdanovskiy"
 __email__ = "vasilyvz@gmail.com"
 __license__ = "MIT"

mcp_security_framework/core/auth_manager.py

@@ -38,6 +38,7 @@ from ..schemas.models import AuthResult, AuthStatus, ValidationResult
 from ..utils.cert_utils import (
     extract_permissions_from_certificate,
     extract_roles_from_certificate,
+    extract_unitid_from_certificate,
     parse_certificate,
     validate_certificate_chain,
 )
@@ -124,20 +125,21 @@ class AuthManager:
         self.logger = logging.getLogger(__name__)
 
         # Initialize storage
-        # Convert API keys from "key": "user" format to "user": "key" format
+        # Store API keys in format: {"api_key": "username"}
         # or handle new format "key": {"username": "user", "roles": ["role1", "role2"]}
         if config.api_keys:
             self._api_keys = {}
             self._api_key_metadata = {}
             for key, value in config.api_keys.items():
                 if isinstance(value, str):
-                    # Old format: "key": "user"
+                    # Old format: "username": "api_key" -> store as {"api_key": "username"}
                     self._api_keys[value] = key
                 elif isinstance(value, dict):
-                    # New format: "key": {"username": "user", "roles": ["role1", "role2"]}
+                    # New format: "username": {"api_key": "key", "roles": ["role1", "role2"]}
+                    api_key = value.get("api_key", key)
                     username = value.get("username", key)
-                    self._api_keys[username] = key
-                    self._api_key_metadata[key] = value
+                    self._api_keys[api_key] = username
+                    self._api_key_metadata[api_key] = value
                 else:
                     self.logger.warning(
                         f"Invalid API key format for key {key}: {value}"
@@ -253,7 +255,7 @@ class AuthManager:
             # Find user by API key
             username = None
             user_roles = []
-            for user, api_key_in_config in self._api_keys.items():
+            for api_key_in_config, user in self._api_keys.items():
                 if api_key_in_config == api_key:
                     username = user
                     # Check if we have metadata for this API key
@@ -625,6 +627,16 @@ class AuthManager:
                 )
                 roles = []
 
+            # Extract unitid from certificate
+            unitid = None
+            try:
+                unitid = extract_unitid_from_certificate(cert_pem)
+            except Exception as e:
+                self.logger.warning(
+                    "Failed to extract unitid from certificate",
+                    extra={"username": username, "error": str(e)},
+                )
+
             # Validate certificate chain if CA is configured
             if self.config.ca_cert_file:
                 try:
@@ -656,6 +668,7 @@ class AuthManager:
                 auth_method="certificate",
                 auth_timestamp=datetime.now(timezone.utc),
                 token_expiry=get_not_valid_after_utc(cert),
+                unitid=unitid,
             )
 
             self.logger.info(
@@ -861,7 +874,7 @@ class AuthManager:
             if not validate_api_key_format(api_key):
                 return False
 
-            self._api_keys[username] = api_key
+            self._api_keys[api_key] = username
 
             self.logger.info("API key added for user", extra={"username": username})
 
@@ -884,8 +897,15 @@ class AuthManager:
             bool: True if API key was removed successfully, False otherwise
         """
         try:
-            if username in self._api_keys:
-                del self._api_keys[username]
+            # Find API key for the username and remove it
+            api_key_to_remove = None
+            for api_key, user in self._api_keys.items():
+                if user == username:
+                    api_key_to_remove = api_key
+                    break
+            
+            if api_key_to_remove:
+                del self._api_keys[api_key_to_remove]
 
                 self.logger.info(
                     "API key removed for user", extra={"username": username}

mcp_security_framework/core/cert_manager.py

@@ -29,6 +29,7 @@ License: MIT
 
 import logging
 import os
+import uuid
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
@@ -54,6 +55,7 @@ from mcp_security_framework.schemas.models import (
 from mcp_security_framework.utils.cert_utils import (
     extract_permissions_from_certificate,
     extract_roles_from_certificate,
+    extract_unitid_from_certificate,
     get_certificate_expiry,
     get_certificate_serial_number,
     get_crl_info,
@@ -254,6 +256,14 @@ class CertificateManager:
                 x509.BasicConstraints(ca=True, path_length=None), critical=True
             )
 
+            # Add unitid extension if provided
+            if ca_config.unitid:
+                unitid_extension = x509.UnrecognizedExtension(
+                    oid=x509.ObjectIdentifier("1.3.6.1.4.1.99999.1.3"),
+                    value=ca_config.unitid.encode(),
+                )
+                builder = builder.add_extension(unitid_extension, critical=False)
+
             builder = builder.add_extension(
                 x509.KeyUsage(
                     digital_signature=True,
@@ -325,6 +335,7 @@ class CertificateManager:
                 not_after=get_not_valid_after_utc(certificate),
                 certificate_type=CertificateType.ROOT_CA,
                 key_size=ca_config.key_size,
+                unitid=ca_config.unitid,
             )
 
             self.logger.info(
@@ -766,6 +777,14 @@ class CertificateManager:
                 )
                 builder = builder.add_extension(permissions_extension, critical=False)
 
+            # Add unitid extension if provided
+            if client_config.unitid:
+                unitid_extension = x509.UnrecognizedExtension(
+                    oid=x509.ObjectIdentifier("1.3.6.1.4.1.99999.1.3"),
+                    value=client_config.unitid.encode(),
+                )
+                builder = builder.add_extension(unitid_extension, critical=False)
+
             # Create certificate
             certificate = builder.sign(ca_key, hashes.SHA256())
 
@@ -816,6 +835,7 @@ class CertificateManager:
                 not_after=get_not_valid_after_utc(certificate),
                 certificate_type=CertificateType.CLIENT,
                 key_size=client_config.key_size,
+                unitid=client_config.unitid,
             )
 
             self.logger.info(

mcp_security_framework/schemas/config.py

@@ -33,6 +33,7 @@ License: MIT
 from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
+import uuid
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 from pydantic.types import SecretStr
@@ -599,6 +600,21 @@ class CAConfig(BaseModel):
     hash_algorithm: str = Field(
         default="sha256", description="Hash algorithm for signing"
     )
+    unitid: Optional[str] = Field(
+        default=None, description="Unique unit identifier (UUID4) for the certificate"
+    )
+
+    @field_validator("unitid")
+    @classmethod
+    def validate_unitid(cls, v):
+        """Validate unitid format."""
+        if v is not None:
+            try:
+                # Validate UUID4 format
+                uuid.UUID(v, version=4)
+            except ValueError:
+                raise ValueError("unitid must be a valid UUID4 string")
+        return v
 
 
 class IntermediateCAConfig(CAConfig):
@@ -668,6 +684,21 @@ class ClientCertConfig(BaseModel):
     )
     ca_cert_path: str = Field(..., description="Path to signing CA certificate")
     ca_key_path: str = Field(..., description="Path to signing CA private key")
+    unitid: Optional[str] = Field(
+        default=None, description="Unique unit identifier (UUID4) for the certificate"
+    )
+
+    @field_validator("unitid")
+    @classmethod
+    def validate_unitid(cls, v):
+        """Validate unitid format."""
+        if v is not None:
+            try:
+                # Validate UUID4 format
+                uuid.UUID(v, version=4)
+            except ValueError:
+                raise ValueError("unitid must be a valid UUID4 string")
+        return v
 
 
 class ServerCertConfig(ClientCertConfig):

mcp_security_framework/schemas/models.py

@@ -38,6 +38,7 @@ License: MIT
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, TypeAlias
+import uuid
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
@@ -140,6 +141,9 @@ class AuthResult(BaseModel):
     metadata: Dict[str, Any] = Field(
         default_factory=dict, description="Additional authentication metadata"
     )
+    unitid: Optional[str] = Field(
+        default=None, description="Unique unit identifier (UUID4) from certificate"
+    )
 
     @field_validator("username")
     @classmethod
@@ -149,6 +153,18 @@ class AuthResult(BaseModel):
             raise ValueError("Username cannot be empty")
         return v
 
+    @field_validator("unitid")
+    @classmethod
+    def validate_unitid(cls, v):
+        """Validate unitid format."""
+        if v is not None:
+            try:
+                # Validate UUID4 format
+                uuid.UUID(v, version=4)
+            except ValueError:
+                raise ValueError("unitid must be a valid UUID4 string")
+        return v
+
     @model_validator(mode="after")
     def validate_auth_result(self):
         """Validate authentication result consistency."""
@@ -309,6 +325,9 @@ class CertificateInfo(BaseModel):
     fingerprint_sha256: Optional[str] = Field(
         default=None, description="SHA256 fingerprint"
     )
+    unitid: Optional[str] = Field(
+        default=None, description="Unique unit identifier (UUID4) for the certificate"
+    )
 
     @field_validator("key_size")
     @classmethod
@@ -318,6 +337,18 @@ class CertificateInfo(BaseModel):
             raise ValueError("Key size must be between 512 and 8192 bits")
         return v
 
+    @field_validator("unitid")
+    @classmethod
+    def validate_unitid(cls, v):
+        """Validate unitid format."""
+        if v is not None:
+            try:
+                # Validate UUID4 format
+                uuid.UUID(v, version=4)
+            except ValueError:
+                raise ValueError("unitid must be a valid UUID4 string")
+        return v
+
     @property
     def is_expired(self) -> bool:
         """Check if certificate is expired."""
@@ -424,6 +455,9 @@ class CertificatePair(BaseModel):
     metadata: Dict[str, Any] = Field(
         default_factory=dict, description="Additional certificate metadata"
     )
+    unitid: Optional[str] = Field(
+        default=None, description="Unique unit identifier (UUID4) for the certificate"
+    )
 
     @field_validator("certificate_pem")
     @classmethod
@@ -449,6 +483,18 @@ class CertificatePair(BaseModel):
             raise ValueError("Invalid private key PEM format")
         return v
 
+    @field_validator("unitid")
+    @classmethod
+    def validate_unitid(cls, v):
+        """Validate unitid format."""
+        if v is not None:
+            try:
+                # Validate UUID4 format
+                uuid.UUID(v, version=4)
+            except ValueError:
+                raise ValueError("unitid must be a valid UUID4 string")
+        return v
+
     @property
     def is_expired(self) -> bool:
         """Check if certificate is expired."""

mcp_security_framework/utils/cert_utils.py

@@ -160,6 +160,15 @@ def extract_certificate_info(cert_data: Union[str, bytes, Path]) -> Dict:
         if country:
             info["country"] = country[0].value
 
+        # Extract unitid
+        try:
+            unitid = extract_unitid_from_certificate(cert_data)
+            if unitid:
+                info["unitid"] = unitid
+        except Exception:
+            # If unitid extraction fails, continue without it
+            pass
+
         return info
     except Exception as e:
         raise CertificateError(f"Certificate information extraction failed: {str(e)}")
@@ -275,6 +284,51 @@ def extract_permissions_from_certificate(
         raise CertificateError(f"Permission extraction failed: {str(e)}")
 
 
+def extract_unitid_from_certificate(
+    cert_data: Union[str, bytes, Path],
+) -> Optional[str]:
+    """
+    Extract unitid from certificate extensions.
+
+    Args:
+        cert_data: Certificate data as string, bytes, or file path
+
+    Returns:
+        Unit ID (UUID4) found in certificate, or None if not found
+
+    Raises:
+        CertificateError: If unitid extraction fails
+    """
+    try:
+        cert = parse_certificate(cert_data)
+        unitid = None
+
+        # Check for custom extension with unitid
+        try:
+            unitid_extension = cert.extensions.get_extension_for_oid(
+                x509.ObjectIdentifier("1.3.6.1.4.1.99999.1.3")  # Custom unitid OID
+            )
+            if unitid_extension:
+                unitid_data = unitid_extension.value.value
+                if isinstance(unitid_data, bytes):
+                    unitid_str = unitid_data.decode("utf-8")
+                    unitid = unitid_str.strip()
+                    
+                    # Validate UUID4 format
+                    try:
+                        import uuid
+                        uuid.UUID(unitid, version=4)
+                    except ValueError:
+                        # Invalid UUID4 format, return None
+                        unitid = None
+        except x509.extensions.ExtensionNotFound:
+            pass
+
+        return unitid
+    except Exception as e:
+        raise CertificateError(f"Unitid extraction failed: {str(e)}")
+
+
 def validate_certificate_chain(
     cert_data: Union[str, bytes, Path],
     ca_cert_data: Union[str, bytes, Path, List[Union[str, bytes, Path]]],

{mcp_security_framework-1.2.0.dist-info → mcp_security_framework-1.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mcp-security-framework
-Version: 1.2.0
+Version: 1.2.1
 Summary: Universal security framework for microservices with SSL/TLS, authentication, authorization, and rate limiting. Requires cryptography>=42.0.0 for certificate operations.
 Author-email: Vasiliy Zdanovskiy <vasilyvz@gmail.com>
 Maintainer-email: Vasiliy Zdanovskiy <vasilyvz@gmail.com>

{mcp_security_framework-1.2.0.dist-info → mcp_security_framework-1.2.1.dist-info}/RECORD

@@ -1,11 +1,11 @@
-mcp_security_framework/__init__.py,sha256=TM8y71Navd_6woEab2cO07MefRsL0tRBItEYfVO7DS8,3172
+mcp_security_framework/__init__.py,sha256=wRlT28hRsnSZCUVNGtU-KggRk66KzSQl4agKOxE5ZCA,3172
 mcp_security_framework/constants.py,sha256=k7NMSrgc83Cci8aoilybQxdC7jir7J-mVFE_EpqVrDk,5307
 mcp_security_framework/cli/__init__.py,sha256=plpWdiWMp2dcLvUuGwXynRg5CDjz8YKnNTBn7lcta08,369
 mcp_security_framework/cli/cert_cli.py,sha256=LdZ3SYKM3e3dP5LsVR5Y0OENtlG0ENu64aHefHjuiN8,23818
 mcp_security_framework/cli/security_cli.py,sha256=Thine_Zzfesz7j29y2k_XZFYUK5YSrhCc6w2FilgEiE,28486
 mcp_security_framework/core/__init__.py,sha256=LiX8_M5qWiTXccJFjSLxup9emhklp-poq57SvznsKEg,1729
-mcp_security_framework/core/auth_manager.py,sha256=6k-Bv7-P7K6TV1KDIzoJGTKYbJmrEUX8uCoZXrNT8Q4,39065
-mcp_security_framework/core/cert_manager.py,sha256=FHwcUAKkTuCRcDANiEBSC3ggHYplDpKR_26qoZKWAdw,88001
+mcp_security_framework/core/auth_manager.py,sha256=GqGAW83Qg1_z2HJ0-FEVTmlli_DBSOPOap2jJMEU1_k,39882
+mcp_security_framework/core/cert_manager.py,sha256=RJgZxElPgMV15Lj_N2uCyO7Ofb9z-MclM1hf8wzFwSc,88882
 mcp_security_framework/core/permission_manager.py,sha256=SADS_oXpwp9MhXHKJMCsvjEq8KWcz7vPYL05Yr-zfio,26478
 mcp_security_framework/core/rate_limiter.py,sha256=6qjVBxK2YHouSxQuCcbr0PBpRqA5toQss_Ce178RElY,20682
 mcp_security_framework/core/security_manager.py,sha256=mAF-5znqxin-MSSgXISB7t1kTkqHltEqGzzmlLAhRGs,37766
@@ -29,12 +29,12 @@ mcp_security_framework/middleware/mtls_middleware.py,sha256=WSyWIk1fCN96hkofODKj
 mcp_security_framework/middleware/rate_limit_middleware.py,sha256=deCwwigI0Pt7pBUnk2jDurI9ZyjujWTsexEWWndXm3g,13177
 mcp_security_framework/middleware/security_middleware.py,sha256=PQ251Fr2UrYVPgGfhXq6QJyqK2tRk0WCIg9_FBvfVkg,16844
 mcp_security_framework/schemas/__init__.py,sha256=lefkbRlbj2ICfasSj51MQ04o3z1YycnbnknSJCFfXbU,2590
-mcp_security_framework/schemas/config.py,sha256=kpJJBupR4ZpQyouZyxiZqm7U7AHfLoGncuRltpo_QEM,25825
-mcp_security_framework/schemas/models.py,sha256=n-Ug8O1cpMeA0mIdOd9h1i3kkBZOJsCQMxj3IcNpBFY,26957
+mcp_security_framework/schemas/config.py,sha256=SUUrpOz9ZTIhZG9Fu2Gsz86yxoGbUt00Qk-Qk_GaTus,26819
+mcp_security_framework/schemas/models.py,sha256=Izjy3I55zjMVLsVZpXZ0M4aK3SCks9sC2U1cbxrXYeI,28439
 mcp_security_framework/schemas/responses.py,sha256=nVXaqF5GTSprXTa_wiUEu38nvSw9WAXtKViAJNbO-Xg,23206
 mcp_security_framework/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mcp_security_framework/utils/__init__.py,sha256=wwwdmQYHTSz0Puvs9FD6aIKmWp3NFARe3JPWNH-b_wk,3098
-mcp_security_framework/utils/cert_utils.py,sha256=roxpUa2CqFHMLvkm8618Epac-U6_xUlNuV8bhEZoU9E,25865
+mcp_security_framework/utils/cert_utils.py,sha256=4AHNrfL0Oqp354aEq5H2_0XLT7v3mNOnTVDV-DLaJyI,27585
 mcp_security_framework/utils/crypto_utils.py,sha256=OH2V7_C3FjStxFTIXMUPfNXZuWG2-QjgoBrIH4Lv4p0,12392
 mcp_security_framework/utils/datetime_compat.py,sha256=ool-xs-EevhuYygdzhiAenLAacLuZwGwjPkF43i-9gg,3859
 mcp_security_framework/utils/validation_utils.py,sha256=e9BX3kw9gdXSmFsc7lmG-qnzSlK0-Ynn7Xs4uKHquF4,16279
@@ -44,7 +44,7 @@ tests/test_cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/test_cli/test_cert_cli.py,sha256=Rm7z-20VAvnmYKY3sgxS-qVNks1vbniQJSpSxjsx_wo,14677
 tests/test_cli/test_security_cli.py,sha256=Bpd31IPJSUl_V1Xzy74ZCOvQpwlbj8Da83C46T8Jewg,25569
 tests/test_core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/test_core/test_auth_manager.py,sha256=tyJUe3yBP95SFwbhSr_IloKhEQw0BuEdvjAC1Hnwu4s,22540
+tests/test_core/test_auth_manager.py,sha256=7Z2DLfJLqKtiwX5Q-lR85hN6NxHbE2Q_FT7IsoyKPQk,22568
 tests/test_core/test_cert_manager.py,sha256=4YMAkRedkAZW3PEZYEbo1PyrzMntUrKfl7arPsHXDCE,36356
 tests/test_core/test_permission_manager.py,sha256=0XeghWXZqVpKyyRuhuDu1dkLUSwuZaFWkRQxQhkkFVI,14966
 tests/test_core/test_rate_limiter.py,sha256=YzzlhlxZm-A7YGMiIV8LXDA0zmb_6uRF9GRx9s21Q0U,22544
@@ -76,9 +76,10 @@ tests/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 tests/test_utils/test_cert_utils.py,sha256=yZGHPuJcjgSHFeT7gnMdsw6UYXmlGUiuHkErukOm8II,28238
 tests/test_utils/test_crypto_utils.py,sha256=yEb4hzG6-irj2DPoXY0DUboJfbeR87ussgTuBpxLGz4,20737
 tests/test_utils/test_datetime_compat.py,sha256=n8S4X5HN-_ejSNpgymDXRyZkmxhnyxwwjxFPdX23I40,5656
+tests/test_utils/test_unitid_compat.py,sha256=MWh03A4FwzQyZa20PKHEWz4W03YtARwBOd_1JbABznQ,25544
 tests/test_utils/test_validation_utils.py,sha256=lus_wHJ2WyVnBGQ28S7dSv78uWcCIuLhn5uflJw-uGw,18569
-mcp_security_framework-1.2.0.dist-info/METADATA,sha256=_mkuF75Qp_BjlrBlamvGIcP-7ejSCdJWHwVsvjZw6eQ,11771
-mcp_security_framework-1.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mcp_security_framework-1.2.0.dist-info/entry_points.txt,sha256=qBh92fVDmd1m2f3xeW0hTu3Ksg8QfGJyV8UEkdA2itg,142
-mcp_security_framework-1.2.0.dist-info/top_level.txt,sha256=ifUiGrTDcD574MXSOoAN2rp2wpUvWlb4jD9LTUgDWCA,29
-mcp_security_framework-1.2.0.dist-info/RECORD,,
+mcp_security_framework-1.2.1.dist-info/METADATA,sha256=3HiQhhOc1aqdstjAhfUnT2gRmxg61pYutP_lVkn6cxw,11771
+mcp_security_framework-1.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mcp_security_framework-1.2.1.dist-info/entry_points.txt,sha256=qBh92fVDmd1m2f3xeW0hTu3Ksg8QfGJyV8UEkdA2itg,142
+mcp_security_framework-1.2.1.dist-info/top_level.txt,sha256=ifUiGrTDcD574MXSOoAN2rp2wpUvWlb4jD9LTUgDWCA,29
+mcp_security_framework-1.2.1.dist-info/RECORD,,

tests/test_core/test_auth_manager.py

@@ -46,7 +46,7 @@ class TestAuthManager:
 
         # Create test configuration
         self.auth_config = AuthConfig(
-            api_keys={"test_api_key_123": "test_user"},
+            api_keys={"test_user": "test_api_key_123"},
             jwt_secret="test_jwt_secret_key_32_chars_long",
             jwt_expiry_hours=1,
             ca_cert_file=None,
@@ -65,7 +65,7 @@ class TestAuthManager:
         """Test successful AuthManager initialization."""
         assert self.auth_manager.config == self.auth_config
         assert self.auth_manager.permission_manager == self.mock_permission_manager
-        assert self.auth_manager._api_keys == {"test_user": "test_api_key_123"}
+        assert self.auth_manager._api_keys == {"test_api_key_123": "test_user"}
         assert self.auth_manager._jwt_secret == "test_jwt_secret_key_32_chars_long"
         assert isinstance(self.auth_manager._token_cache, dict)
         assert isinstance(self.auth_manager._session_store, dict)
@@ -429,8 +429,8 @@ class TestAuthManager:
         success = self.auth_manager.add_api_key("new_user", "new_api_key_456789")
 
         assert success is True
-        assert "new_user" in self.auth_manager._api_keys
-        assert self.auth_manager._api_keys["new_user"] == "new_api_key_456789"
+        assert "new_api_key_456789" in self.auth_manager._api_keys
+        assert self.auth_manager._api_keys["new_api_key_456789"] == "new_user"
 
     def test_add_api_key_invalid_input(self):
         """Test API key addition with invalid input."""
@@ -456,12 +456,12 @@ class TestAuthManager:
         """Test successful API key removal."""
         # Add a key first
         self.auth_manager.add_api_key("temp_user", "temp_key_123456789")
-        assert "temp_user" in self.auth_manager._api_keys
+        assert "temp_key_123456789" in self.auth_manager._api_keys
 
         # Remove the key
         success = self.auth_manager.remove_api_key("temp_user")
         assert success is True
-        assert "temp_user" not in self.auth_manager._api_keys
+        assert "temp_key_123456789" not in self.auth_manager._api_keys
 
     def test_remove_api_key_not_found(self):
         """Test API key removal for non-existent user."""

tests/test_utils/test_unitid_compat.py

@@ -0,0 +1,550 @@
+"""
+Tests for UnitID Compatibility Module
+
+This module contains tests for the unitid functionality that
+handles unique unit identifiers in certificates.
+
+Test Coverage:
+- UnitID extraction from certificates
+- UnitID validation
+- UnitID integration with certificate creation
+- UnitID integration with authentication
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import pytest
+import uuid
+from unittest.mock import Mock, patch
+from cryptography import x509
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+from mcp_security_framework.utils.cert_utils import extract_unitid_from_certificate
+
+
+class TestUnitIDCompatibility:
+    """Test suite for unitid compatibility functions."""
+
+    def test_extract_unitid_from_certificate_with_unitid_extension(self):
+        """Test extracting unitid from certificate with unitid extension."""
+        # Create a mock certificate with unitid extension
+        mock_cert = Mock()
+        mock_extension = Mock()
+        mock_extension.value.value = b"550e8400-e29b-41d4-a716-446655440000"
+        mock_cert.extensions.get_extension_for_oid.return_value = mock_extension
+        
+        # Mock the parse_certificate function to return our mock certificate
+        with patch('mcp_security_framework.utils.cert_utils.parse_certificate') as mock_parse:
+            mock_parse.return_value = mock_cert
+            
+            result = extract_unitid_from_certificate("mock_cert_data")
+            
+            assert result == "550e8400-e29b-41d4-a716-446655440000"
+
+    def test_extract_unitid_from_certificate_without_unitid_extension(self):
+        """Test extracting unitid from certificate without unitid extension."""
+        # Create a mock certificate without unitid extension
+        mock_cert = Mock()
+        mock_cert.extensions.get_extension_for_oid.side_effect = x509.ExtensionNotFound(
+            "Extension not found", "1.3.6.1.4.1.99999.1.3"
+        )
+        
+        # Mock the parse_certificate function to return our mock certificate
+        with patch('mcp_security_framework.utils.cert_utils.parse_certificate') as mock_parse:
+            mock_parse.return_value = mock_cert
+            
+            result = extract_unitid_from_certificate("mock_cert_data")
+            
+            assert result is None
+
+    def test_extract_unitid_from_certificate_with_invalid_unitid(self):
+        """Test extracting unitid from certificate with invalid unitid format."""
+        # Create a mock certificate with invalid unitid
+        mock_cert = Mock()
+        mock_extension = Mock()
+        mock_extension.value.value = b"invalid-unitid"
+        mock_cert.extensions.get_extension_for_oid.return_value = mock_extension
+        
+        # Mock the parse_certificate function to return our mock certificate
+        with patch('mcp_security_framework.utils.cert_utils.parse_certificate') as mock_parse:
+            mock_parse.return_value = mock_cert
+            
+            result = extract_unitid_from_certificate("mock_cert_data")
+            
+            # Should return None for invalid unitid
+            assert result is None
+
+    def test_extract_unitid_from_certificate_with_real_certificate(self):
+        """Test extracting unitid from a real certificate."""
+        # Generate a valid UUID4
+        test_unitid = str(uuid.uuid4())
+        
+        # Create a real certificate with unitid extension
+        private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=2048,
+        )
+
+        subject = issuer = x509.Name([
+            x509.NameAttribute(x509.NameOID.COMMON_NAME, 'Test UnitID'),
+        ])
+
+        # Create certificate builder
+        from datetime import datetime, timezone, timedelta
+        builder = x509.CertificateBuilder()
+        builder = builder.subject_name(subject)
+        builder = builder.issuer_name(issuer)
+        builder = builder.public_key(private_key.public_key())
+        builder = builder.serial_number(x509.random_serial_number())
+        builder = builder.not_valid_before(datetime.now(timezone.utc))
+        builder = builder.not_valid_after(
+            datetime.now(timezone.utc) + timedelta(days=10)
+        )
+
+        # Add unitid extension
+        unitid_extension = x509.UnrecognizedExtension(
+            oid=x509.ObjectIdentifier("1.3.6.1.4.1.99999.1.3"),
+            value=test_unitid.encode(),
+        )
+        builder = builder.add_extension(unitid_extension, critical=False)
+
+        # Sign the certificate
+        certificate = builder.sign(private_key, hashes.SHA256())
+
+        # Test extraction
+        cert_pem = certificate.public_bytes(serialization.Encoding.PEM)
+        result = extract_unitid_from_certificate(cert_pem)
+        
+        assert result == test_unitid
+
+    def test_unitid_validation_in_certificate_info(self):
+        """Test unitid validation in CertificateInfo model."""
+        from mcp_security_framework.schemas.models import CertificateInfo, CertificateType
+        from datetime import datetime, timezone
+
+        # Test with valid UUID4
+        valid_unitid = str(uuid.uuid4())
+        cert_info = CertificateInfo(
+            subject={"CN": "Test"},
+            issuer={"CN": "Test CA"},
+            serial_number="123456789",
+            not_before=datetime.now(timezone.utc),
+            not_after=datetime.now(timezone.utc),
+            certificate_type=CertificateType.CLIENT,
+            key_size=2048,
+            signature_algorithm="sha256",
+            unitid=valid_unitid
+        )
+        
+        assert cert_info.unitid == valid_unitid
+
+        # Test with invalid UUID4
+        with pytest.raises(ValueError, match="unitid must be a valid UUID4 string"):
+            CertificateInfo(
+                subject={"CN": "Test"},
+                issuer={"CN": "Test CA"},
+                serial_number="123456789",
+                not_before=datetime.now(timezone.utc),
+                not_after=datetime.now(timezone.utc),
+                certificate_type=CertificateType.CLIENT,
+                key_size=2048,
+                signature_algorithm="sha256",
+                unitid="invalid-unitid"
+            )
+
+    def test_unitid_validation_in_certificate_pair(self):
+        """Test unitid validation in CertificatePair model."""
+        from mcp_security_framework.schemas.models import CertificatePair, CertificateType
+        from datetime import datetime, timezone
+
+        # Test with valid UUID4
+        valid_unitid = str(uuid.uuid4())
+        cert_pair = CertificatePair(
+            certificate_path="/path/to/cert.crt",
+            private_key_path="/path/to/key.key",
+            certificate_pem="-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----",
+            private_key_pem="-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----",
+            serial_number="123456789",
+            common_name="Test",
+            organization="Test Org",
+            not_before=datetime.now(timezone.utc),
+            not_after=datetime.now(timezone.utc),
+            certificate_type=CertificateType.CLIENT,
+            key_size=2048,
+            unitid=valid_unitid
+        )
+        
+        assert cert_pair.unitid == valid_unitid
+
+        # Test with invalid UUID4
+        with pytest.raises(ValueError, match="unitid must be a valid UUID4 string"):
+            CertificatePair(
+                certificate_path="/path/to/cert.crt",
+                private_key_path="/path/to/key.key",
+                certificate_pem="-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----",
+                private_key_pem="-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----",
+                serial_number="123456789",
+                common_name="Test",
+                organization="Test Org",
+                not_before=datetime.now(timezone.utc),
+                not_after=datetime.now(timezone.utc),
+                certificate_type=CertificateType.CLIENT,
+                key_size=2048,
+                unitid="invalid-unitid"
+            )
+
+    def test_unitid_validation_in_auth_result(self):
+        """Test unitid validation in AuthResult model."""
+        from mcp_security_framework.schemas.models import AuthResult, AuthStatus
+        from datetime import datetime, timezone
+
+        # Test with valid UUID4
+        valid_unitid = str(uuid.uuid4())
+        auth_result = AuthResult(
+            is_valid=True,
+            status=AuthStatus.SUCCESS,
+            username="test_user",
+            unitid=valid_unitid
+        )
+        
+        assert auth_result.unitid == valid_unitid
+
+        # Test with invalid UUID4
+        with pytest.raises(ValueError, match="unitid must be a valid UUID4 string"):
+            AuthResult(
+                is_valid=True,
+                status=AuthStatus.SUCCESS,
+                username="test_user",
+                unitid="invalid-unitid"
+            )
+
+    def test_unitid_in_certificate_creation_config(self):
+        """Test unitid in certificate creation configuration."""
+        from mcp_security_framework.schemas.config import CAConfig, ClientCertConfig
+
+        # Test with valid UUID4
+        valid_unitid = str(uuid.uuid4())
+        
+        ca_config = CAConfig(
+            common_name="Test CA",
+            organization="Test Org",
+            country="US",
+            unitid=valid_unitid
+        )
+        assert ca_config.unitid == valid_unitid
+
+        client_config = ClientCertConfig(
+            common_name="test.example.com",
+            organization="Test Org",
+            country="US",
+            ca_cert_path="/path/to/ca.crt",
+            ca_key_path="/path/to/ca.key",
+            unitid=valid_unitid
+        )
+        assert client_config.unitid == valid_unitid
+
+        # Test with invalid UUID4
+        with pytest.raises(ValueError, match="unitid must be a valid UUID4 string"):
+            CAConfig(
+                common_name="Test CA",
+                organization="Test Org",
+                country="US",
+                unitid="invalid-unitid"
+            )
+
+    def test_unitid_integration_with_certificate_creation(self):
+        """Test unitid integration with certificate creation."""
+        from mcp_security_framework.core.cert_manager import CertificateManager
+        from mcp_security_framework.schemas.config import CertificateConfig, CAConfig
+        from datetime import datetime, timezone
+        import tempfile
+        import os
+
+        # Create temporary directory for certificates
+        with tempfile.TemporaryDirectory() as temp_dir:
+            cert_config = CertificateConfig(
+                enabled=True,
+                ca_cert_path=os.path.join(temp_dir, "ca.crt"),
+                ca_key_path=os.path.join(temp_dir, "ca.key"),
+                cert_storage_path=temp_dir,
+                key_storage_path=temp_dir
+            )
+
+            # Mock the configuration validation to avoid file system checks
+            with patch.object(CertificateManager, '_validate_configuration'):
+                cert_manager = CertificateManager(cert_config)
+
+            # Test unitid in CA creation
+            valid_unitid = str(uuid.uuid4())
+            ca_config = CAConfig(
+                common_name="Test CA",
+                organization="Test Org",
+                country="US",
+                unitid=valid_unitid
+            )
+
+            # Mock the certificate creation to avoid actual file operations
+            with patch.object(cert_manager, '_validate_configuration'):
+                with patch('builtins.open', create=True):
+                    with patch('os.makedirs'):
+                        with patch('os.chmod'):
+                            # Mock the certificate building process
+                            mock_cert = Mock()
+                            mock_cert.serial_number = 123456789
+                            mock_cert.not_valid_before = datetime.now(timezone.utc)
+                            mock_cert.not_valid_after = datetime.now(timezone.utc)
+                            mock_cert.public_bytes.return_value = b"-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----"
+                            
+                            with patch('mcp_security_framework.core.cert_manager.rsa.generate_private_key') as mock_key:
+                                mock_private_key = Mock()
+                                mock_private_key.public_key.return_value = Mock()
+                                mock_private_key.private_bytes.return_value = b"-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----"
+                                mock_key.return_value = mock_private_key
+                                
+                                with patch('mcp_security_framework.core.cert_manager.x509.CertificateBuilder') as mock_builder:
+                                    mock_builder_instance = Mock()
+                                    mock_builder_instance.subject_name.return_value = mock_builder_instance
+                                    mock_builder_instance.issuer_name.return_value = mock_builder_instance
+                                    mock_builder_instance.public_key.return_value = mock_builder_instance
+                                    mock_builder_instance.serial_number.return_value = mock_builder_instance
+                                    mock_builder_instance.not_valid_before.return_value = mock_builder_instance
+                                    mock_builder_instance.not_valid_after.return_value = mock_builder_instance
+                                    mock_builder_instance.add_extension.return_value = mock_builder_instance
+                                    mock_builder_instance.sign.return_value = mock_cert
+                                    mock_builder.return_value = mock_builder_instance
+                                    
+                                    # Mock the datetime compatibility functions
+                                    with patch('mcp_security_framework.core.cert_manager.get_not_valid_before_utc') as mock_before:
+                                        with patch('mcp_security_framework.core.cert_manager.get_not_valid_after_utc') as mock_after:
+                                            mock_before.return_value = datetime.now(timezone.utc)
+                                            mock_after.return_value = datetime.now(timezone.utc)
+                                            
+                                            result = cert_manager.create_root_ca(ca_config)
+                                            
+                                            # Verify unitid was added to the certificate pair
+                                            assert result.unitid == valid_unitid
+
+    def test_unitid_integration_with_authentication(self):
+        """Test unitid integration with authentication."""
+        from mcp_security_framework.core.auth_manager import AuthManager
+        from mcp_security_framework.schemas.config import AuthConfig
+        from mcp_security_framework.core.permission_manager import PermissionManager
+        from mcp_security_framework.schemas.config import PermissionConfig
+        from unittest.mock import Mock
+        import tempfile
+        import json
+        import os
+
+        # Create test configuration
+        auth_config = AuthConfig(
+            enabled=True,
+            api_keys={"test_key": "test_user"},
+            jwt_secret="test_secret_key_for_jwt_signing_12345"
+        )
+
+        # Create temporary roles file
+        with tempfile.TemporaryDirectory() as temp_dir:
+            roles_file = os.path.join(temp_dir, "test_roles.json")
+            with open(roles_file, "w") as f:
+                json.dump({"roles": {}}, f)
+
+            perm_config = PermissionConfig(
+                enabled=True,
+                roles_file=roles_file
+            )
+
+            permission_manager = PermissionManager(perm_config)
+            auth_manager = AuthManager(auth_config, permission_manager)
+
+        # Test unitid in certificate authentication
+        valid_unitid = str(uuid.uuid4())
+        
+        # Mock certificate with unitid
+        mock_cert_pem = "-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----"
+        
+        # Mock the authenticate_certificate method to return a successful result with unitid
+        from mcp_security_framework.schemas.models import AuthResult, AuthStatus
+        from datetime import datetime, timezone, timedelta
+        
+        expected_result = AuthResult(
+            is_valid=True,
+            status=AuthStatus.SUCCESS,
+            username="test_user",
+            roles=[],
+            auth_method="certificate",
+            auth_timestamp=datetime.now(timezone.utc),
+            token_expiry=datetime.now(timezone.utc) + timedelta(days=30),
+            unitid=valid_unitid,
+        )
+        
+        with patch.object(auth_manager, 'authenticate_certificate', return_value=expected_result):
+            result = auth_manager.authenticate_certificate(mock_cert_pem)
+            
+            # Verify unitid was included in authentication result
+            assert result.unitid == valid_unitid
+
+    def test_unitid_optional_in_certificate_creation(self):
+        """Test that unitid is optional in certificate creation."""
+        from mcp_security_framework.schemas.config import CAConfig, ClientCertConfig
+
+        # Test CA config without unitid
+        ca_config = CAConfig(
+            common_name="Test CA",
+            organization="Test Org",
+            country="US"
+            # unitid not specified
+        )
+        assert ca_config.unitid is None
+
+        # Test client config without unitid
+        client_config = ClientCertConfig(
+            common_name="test.example.com",
+            organization="Test Org",
+            country="US",
+            ca_cert_path="/path/to/ca.crt",
+            ca_key_path="/path/to/ca.key"
+            # unitid not specified
+        )
+        assert client_config.unitid is None
+
+    def test_unitid_optional_in_certificate_info(self):
+        """Test that unitid is optional in CertificateInfo."""
+        from mcp_security_framework.schemas.models import CertificateInfo, CertificateType
+        from datetime import datetime, timezone
+
+        # Test CertificateInfo without unitid
+        cert_info = CertificateInfo(
+            subject={"CN": "Test"},
+            issuer={"CN": "Test CA"},
+            serial_number="123456789",
+            not_before=datetime.now(timezone.utc),
+            not_after=datetime.now(timezone.utc),
+            certificate_type=CertificateType.CLIENT,
+            key_size=2048,
+            signature_algorithm="sha256"
+            # unitid not specified
+        )
+        assert cert_info.unitid is None
+
+    def test_unitid_optional_in_certificate_pair(self):
+        """Test that unitid is optional in CertificatePair."""
+        from mcp_security_framework.schemas.models import CertificatePair, CertificateType
+        from datetime import datetime, timezone
+
+        # Test CertificatePair without unitid
+        cert_pair = CertificatePair(
+            certificate_path="/path/to/cert.crt",
+            private_key_path="/path/to/key.key",
+            certificate_pem="-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----",
+            private_key_pem="-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----",
+            serial_number="123456789",
+            common_name="Test",
+            organization="Test Org",
+            not_before=datetime.now(timezone.utc),
+            not_after=datetime.now(timezone.utc),
+            certificate_type=CertificateType.CLIENT,
+            key_size=2048
+            # unitid not specified
+        )
+        assert cert_pair.unitid is None
+
+    def test_unitid_optional_in_auth_result(self):
+        """Test that unitid is optional in AuthResult."""
+        from mcp_security_framework.schemas.models import AuthResult, AuthStatus
+
+        # Test AuthResult without unitid
+        auth_result = AuthResult(
+            is_valid=True,
+            status=AuthStatus.SUCCESS,
+            username="test_user"
+            # unitid not specified
+        )
+        assert auth_result.unitid is None
+
+    def test_extract_unitid_from_certificate_without_unitid_returns_none(self):
+        """Test that extract_unitid_from_certificate returns None when unitid is not present."""
+        # Create a mock certificate without unitid extension
+        mock_cert = Mock()
+        mock_cert.extensions.get_extension_for_oid.side_effect = x509.ExtensionNotFound(
+            "Extension not found", "1.3.6.1.4.1.99999.1.3"
+        )
+        
+        # Mock the parse_certificate function to return our mock certificate
+        with patch('mcp_security_framework.utils.cert_utils.parse_certificate') as mock_parse:
+            mock_parse.return_value = mock_cert
+            
+            result = extract_unitid_from_certificate("mock_cert_data")
+            
+            # Should return None when unitid is not present
+            assert result is None
+
+    def test_certificate_creation_without_unitid(self):
+        """Test certificate creation when unitid is not specified."""
+        from mcp_security_framework.core.cert_manager import CertificateManager
+        from mcp_security_framework.schemas.config import CertificateConfig, CAConfig
+        from datetime import datetime, timezone
+        import tempfile
+        import os
+
+        # Create temporary directory for certificates
+        with tempfile.TemporaryDirectory() as temp_dir:
+            cert_config = CertificateConfig(
+                enabled=True,
+                ca_cert_path=os.path.join(temp_dir, "ca.crt"),
+                ca_key_path=os.path.join(temp_dir, "ca.key"),
+                cert_storage_path=temp_dir,
+                key_storage_path=temp_dir
+            )
+
+            # Mock the configuration validation to avoid file system checks
+            with patch.object(CertificateManager, '_validate_configuration'):
+                cert_manager = CertificateManager(cert_config)
+
+                # Test CA creation without unitid
+                ca_config = CAConfig(
+                    common_name="Test CA",
+                    organization="Test Org",
+                    country="US"
+                    # unitid not specified
+                )
+
+                # Mock the certificate building process
+                with patch('builtins.open', create=True):
+                    with patch('os.makedirs'):
+                        with patch('os.chmod'):
+                            mock_cert = Mock()
+                            mock_cert.serial_number = 123456789
+                            mock_cert.not_valid_before = datetime.now(timezone.utc)
+                            mock_cert.not_valid_after = datetime.now(timezone.utc)
+                            mock_cert.public_bytes.return_value = b"-----BEGIN CERTIFICATE-----\nMOCK\n-----END CERTIFICATE-----"
+                            
+                            with patch('mcp_security_framework.core.cert_manager.rsa.generate_private_key') as mock_key:
+                                mock_private_key = Mock()
+                                mock_private_key.public_key.return_value = Mock()
+                                mock_private_key.private_bytes.return_value = b"-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----"
+                                mock_key.return_value = mock_private_key
+                                
+                                with patch('mcp_security_framework.core.cert_manager.x509.CertificateBuilder') as mock_builder:
+                                    mock_builder_instance = Mock()
+                                    mock_builder_instance.subject_name.return_value = mock_builder_instance
+                                    mock_builder_instance.issuer_name.return_value = mock_builder_instance
+                                    mock_builder_instance.public_key.return_value = mock_builder_instance
+                                    mock_builder_instance.serial_number.return_value = mock_builder_instance
+                                    mock_builder_instance.not_valid_before.return_value = mock_builder_instance
+                                    mock_builder_instance.not_valid_after.return_value = mock_builder_instance
+                                    mock_builder_instance.add_extension.return_value = mock_builder_instance
+                                    mock_builder_instance.sign.return_value = mock_cert
+                                    mock_builder.return_value = mock_builder_instance
+                                    
+                                    # Mock the datetime compatibility functions
+                                    with patch('mcp_security_framework.core.cert_manager.get_not_valid_before_utc') as mock_before:
+                                        with patch('mcp_security_framework.core.cert_manager.get_not_valid_after_utc') as mock_after:
+                                            mock_before.return_value = datetime.now(timezone.utc)
+                                            mock_after.return_value = datetime.now(timezone.utc)
+                                            
+                                            result = cert_manager.create_root_ca(ca_config)
+                                            
+                                            # Verify unitid is None when not specified
+                                            assert result.unitid is None