feat: implement EntraID auth for Azure with Service Principal, Managed Identity and Default Credential support

vchomakov · vchomakov · commit 23dab1ab6e95 · 2025-10-22T18:20:54.000+03:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -26,6 +26,7 @@ dependencies = [
     "dotenv>=0.9.9",
     "numpy>=2.2.4",
     "click>=8.0.0",
+    "redis-entraid>=1.0.0",
 ]
 
 [project.scripts]
diff --git a/src/common/config.py b/src/common/config.py
@@ -5,6 +5,13 @@
 
 load_dotenv()
 
+# Default values for Entra ID authentication
+DEFAULT_TOKEN_EXPIRATION_REFRESH_RATIO = 0.9
+DEFAULT_LOWER_REFRESH_BOUND_MILLIS = 30000  # 30 seconds
+DEFAULT_TOKEN_REQUEST_EXECUTION_TIMEOUT_MS = 10000  # 10 seconds
+DEFAULT_RETRY_MAX_ATTEMPTS = 3
+DEFAULT_RETRY_DELAY_MS = 100
+
 REDIS_CFG = {
     "host": os.getenv("REDIS_HOST", "127.0.0.1"),
     "port": int(os.getenv("REDIS_PORT", 6379)),
@@ -20,6 +27,55 @@
     "db": int(os.getenv("REDIS_DB", 0)),
 }
 
+# Entra ID Authentication Configuration
+ENTRAID_CFG = {
+    # Authentication flow selection
+    "auth_flow": os.getenv(
+        "REDIS_ENTRAID_AUTH_FLOW", None
+    ),  # service_principal, managed_identity, default_credential
+    # Service Principal Authentication
+    "client_id": os.getenv("REDIS_ENTRAID_CLIENT_ID", None),
+    "client_secret": os.getenv("REDIS_ENTRAID_CLIENT_SECRET", None),
+    "tenant_id": os.getenv("REDIS_ENTRAID_TENANT_ID", None),
+    # Managed Identity Authentication
+    "identity_type": os.getenv(
+        "REDIS_ENTRAID_IDENTITY_TYPE", "system_assigned"
+    ),  # system_assigned, user_assigned
+    "user_assigned_identity_client_id": os.getenv(
+        "REDIS_ENTRAID_USER_ASSIGNED_CLIENT_ID", None
+    ),
+    # Default Azure Credential Authentication
+    "scopes": os.getenv("REDIS_ENTRAID_SCOPES", "https://redis.azure.com/.default"),
+    # Token lifecycle configuration
+    "token_expiration_refresh_ratio": float(
+        os.getenv(
+            "REDIS_ENTRAID_TOKEN_EXPIRATION_REFRESH_RATIO",
+            DEFAULT_TOKEN_EXPIRATION_REFRESH_RATIO,
+        )
+    ),
+    "lower_refresh_bound_millis": int(
+        os.getenv(
+            "REDIS_ENTRAID_LOWER_REFRESH_BOUND_MILLIS",
+            DEFAULT_LOWER_REFRESH_BOUND_MILLIS,
+        )
+    ),
+    "token_request_execution_timeout_ms": int(
+        os.getenv(
+            "REDIS_ENTRAID_TOKEN_REQUEST_EXECUTION_TIMEOUT_MS",
+            DEFAULT_TOKEN_REQUEST_EXECUTION_TIMEOUT_MS,
+        )
+    ),
+    # Retry configuration
+    "retry_max_attempts": int(
+        os.getenv("REDIS_ENTRAID_RETRY_MAX_ATTEMPTS", DEFAULT_RETRY_MAX_ATTEMPTS)
+    ),
+    "retry_delay_ms": int(
+        os.getenv("REDIS_ENTRAID_RETRY_DELAY_MS", DEFAULT_RETRY_DELAY_MS)
+    ),
+    # Resource configuration
+    "resource": os.getenv("REDIS_ENTRAID_RESOURCE", "https://redis.azure.com/"),
+}
+
 
 def parse_redis_uri(uri: str) -> dict:
     """Parse a Redis URI and return connection parameters."""
@@ -99,3 +155,77 @@ def set_redis_config_from_cli(config: dict):
         else:
             # Convert other values to strings
             REDIS_CFG[key] = str(value) if value is not None else None
+
+
+def set_entraid_config_from_cli(config: dict):
+    """Update Entra ID configuration from CLI parameters."""
+    for key, value in config.items():
+        if value is not None:
+            if key in ["token_expiration_refresh_ratio"]:
+                # Keep float values as floats
+                ENTRAID_CFG[key] = float(value)
+            elif key in [
+                "lower_refresh_bound_millis",
+                "token_request_execution_timeout_ms",
+                "retry_max_attempts",
+                "retry_delay_ms",
+            ]:
+                # Keep integer values as integers
+                ENTRAID_CFG[key] = int(value)
+            else:
+                # Convert other values to strings
+                ENTRAID_CFG[key] = str(value)
+
+
+def is_entraid_auth_enabled() -> bool:
+    """Check if Entra ID authentication is enabled."""
+    return ENTRAID_CFG["auth_flow"] is not None
+
+
+def get_entraid_auth_flow() -> str:
+    """Get the configured Entra ID authentication flow."""
+    return ENTRAID_CFG["auth_flow"]
+
+
+def validate_entraid_config() -> tuple[bool, str]:
+    """Validate Entra ID configuration based on the selected auth flow.
+
+    Returns:
+        tuple: (is_valid, error_message)
+    """
+    auth_flow = ENTRAID_CFG["auth_flow"]
+
+    if not auth_flow:
+        return True, ""  # No Entra ID auth configured, which is valid
+
+    if auth_flow == "service_principal":
+        required_fields = ["client_id", "client_secret", "tenant_id"]
+        missing_fields = [field for field in required_fields if not ENTRAID_CFG[field]]
+        if missing_fields:
+            return (
+                False,
+                f"Service principal authentication requires: {', '.join(missing_fields)}",
+            )
+
+    elif auth_flow == "managed_identity":
+        identity_type = ENTRAID_CFG["identity_type"]
+        if (
+            identity_type == "user_assigned"
+            and not ENTRAID_CFG["user_assigned_identity_client_id"]
+        ):
+            return (
+                False,
+                "User-assigned managed identity requires user_assigned_identity_client_id",
+            )
+
+    elif auth_flow == "default_credential":
+        # Default credential doesn't require specific configuration
+        pass
+
+    else:
+        return (
+            False,
+            f"Invalid auth_flow: {auth_flow}. Must be one of: service_principal, managed_identity, default_credential",
+        )
+
+    return True, ""
diff --git a/src/common/connection.py b/src/common/connection.py
@@ -5,7 +5,11 @@
 from redis import Redis
 from redis.cluster import RedisCluster
 
-from src.common.config import REDIS_CFG
+from src.common.config import REDIS_CFG, is_entraid_auth_enabled
+from src.common.entraid_auth import (
+    create_credential_provider,
+    EntraIDAuthenticationError,
+)
 from src.version import __version__
 
 _logger = logging.getLogger(__name__)
@@ -18,6 +22,17 @@ class RedisConnectionManager:
     def get_connection(cls, decode_responses=True) -> Redis:
         if cls._instance is None:
             try:
+                # Create Entra ID credential provider if configured
+                credential_provider = None
+                if is_entraid_auth_enabled():
+                    try:
+                        credential_provider = create_credential_provider()
+                    except EntraIDAuthenticationError as e:
+                        _logger.error(
+                            "Failed to create Entra ID credential provider: %s", e
+                        )
+                        raise
+
                 if REDIS_CFG["cluster_mode"]:
                     redis_class: Type[Union[Redis, RedisCluster]] = (
                         redis.cluster.RedisCluster
@@ -37,6 +52,12 @@ def get_connection(cls, decode_responses=True) -> Redis:
                         "lib_name": f"redis-py(mcp-server_v{__version__})",
                         "max_connections_per_node": 10,
                     }
+
+                    # Add credential provider if available
+                    if credential_provider:
+                        connection_params["credential_provider"] = credential_provider
+                        # Note: Azure Redis Enterprise with EntraID uses plain text connections
+                        # SSL setting is controlled by REDIS_SSL environment variable
                 else:
                     redis_class: Type[Union[Redis, RedisCluster]] = redis.Redis
                     connection_params = {
@@ -56,6 +77,12 @@ def get_connection(cls, decode_responses=True) -> Redis:
                         "max_connections": 10,
                     }
 
+                    # Add credential provider if available
+                    if credential_provider:
+                        connection_params["credential_provider"] = credential_provider
+                        # Note: Azure Redis Enterprise with EntraID uses plain text connections
+                        # SSL setting is controlled by REDIS_SSL environment variable
+
                 cls._instance = redis_class(**connection_params)
 
             except redis.exceptions.ConnectionError:
diff --git a/src/common/entraid_auth.py b/src/common/entraid_auth.py
@@ -0,0 +1,160 @@
+"""
+Entra ID authentication provider factory for Redis MCP Server.
+
+This module provides factory methods to create credential providers for different
+Azure authentication flows based on configuration.
+"""
+
+import logging
+
+from src.common.config import (
+    ENTRAID_CFG,
+    is_entraid_auth_enabled,
+    validate_entraid_config,
+)
+
+_logger = logging.getLogger(__name__)
+
+# Reduce Azure SDK logging verbosity
+logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(logging.WARNING)
+logging.getLogger("azure.identity").setLevel(logging.WARNING)
+logging.getLogger("redis.auth.token_manager").setLevel(logging.WARNING)
+
+# Import redis-entraid components only when needed
+try:
+    from redis_entraid.cred_provider import (
+        create_from_default_azure_credential,
+        create_from_managed_identity,
+        create_from_service_principal,
+        ManagedIdentityType,
+        TokenManagerConfig,
+        RetryPolicy,
+    )
+
+    ENTRAID_AVAILABLE = True
+except ImportError:
+    _logger.warning(
+        "redis-entraid package not available. Entra ID authentication will be disabled."
+    )
+    ENTRAID_AVAILABLE = False
+
+
+class EntraIDAuthenticationError(Exception):
+    """Exception raised for Entra ID authentication configuration errors."""
+
+    pass
+
+
+def create_credential_provider():
+    """
+    Create an Entra ID credential provider based on the current configuration.
+
+    Returns:
+        Credential provider instance or None if Entra ID auth is not configured.
+
+    Raises:
+        EntraIDAuthenticationError: If configuration is invalid or required packages are missing.
+    """
+    if not is_entraid_auth_enabled():
+        return None
+
+    if not ENTRAID_AVAILABLE:
+        raise EntraIDAuthenticationError(
+            "redis-entraid package is required for Entra ID authentication. "
+            "Install it with: pip install redis-entraid"
+        )
+
+    # Validate configuration
+    is_valid, error_message = validate_entraid_config()
+    if not is_valid:
+        raise EntraIDAuthenticationError(
+            f"Invalid Entra ID configuration: {error_message}"
+        )
+
+    auth_flow = ENTRAID_CFG["auth_flow"]
+
+    try:
+        # Create token manager configuration
+        token_manager_config = _create_token_manager_config()
+
+        if auth_flow == "service_principal":
+            return _create_service_principal_provider(token_manager_config)
+        elif auth_flow == "managed_identity":
+            return _create_managed_identity_provider(token_manager_config)
+        elif auth_flow == "default_credential":
+            return _create_default_credential_provider(token_manager_config)
+        else:
+            raise EntraIDAuthenticationError(
+                f"Unsupported authentication flow: {auth_flow}"
+            )
+
+    except Exception as e:
+        _logger.error("Failed to create Entra ID credential provider: %s", e)
+        raise EntraIDAuthenticationError(f"Failed to create credential provider: {e}")
+
+
+def _create_token_manager_config():
+    """Create TokenManagerConfig from current configuration."""
+    retry_policy = RetryPolicy(
+        max_attempts=ENTRAID_CFG["retry_max_attempts"],
+        delay_in_ms=ENTRAID_CFG["retry_delay_ms"],
+    )
+
+    return TokenManagerConfig(
+        expiration_refresh_ratio=ENTRAID_CFG["token_expiration_refresh_ratio"],
+        lower_refresh_bound_millis=ENTRAID_CFG["lower_refresh_bound_millis"],
+        token_request_execution_timeout_in_ms=ENTRAID_CFG[
+            "token_request_execution_timeout_ms"
+        ],
+        retry_policy=retry_policy,
+    )
+
+
+def _create_service_principal_provider(token_manager_config):
+    """Create service principal credential provider."""
+
+    return create_from_service_principal(
+        client_id=ENTRAID_CFG["client_id"],
+        client_credential=ENTRAID_CFG["client_secret"],
+        tenant_id=ENTRAID_CFG["tenant_id"],
+        token_manager_config=token_manager_config,
+    )
+
+
+def _create_managed_identity_provider(token_manager_config):
+    """Create managed identity credential provider."""
+    identity_type_str = ENTRAID_CFG["identity_type"]
+
+    if identity_type_str == "system_assigned":
+        identity_type = ManagedIdentityType.SYSTEM_ASSIGNED
+
+        return create_from_managed_identity(
+            identity_type=identity_type,
+            resource=ENTRAID_CFG["resource"],
+            token_manager_config=token_manager_config,
+        )
+
+    elif identity_type_str == "user_assigned":
+        identity_type = ManagedIdentityType.USER_ASSIGNED
+
+        return create_from_managed_identity(
+            identity_type=identity_type,
+            resource=ENTRAID_CFG["resource"],
+            client_id=ENTRAID_CFG["user_assigned_identity_client_id"],
+            token_manager_config=token_manager_config,
+        )
+
+    else:
+        raise EntraIDAuthenticationError(f"Invalid identity type: {identity_type_str}")
+
+
+def _create_default_credential_provider(token_manager_config):
+    """Create default Azure credential provider."""
+
+    # Parse scopes from configuration
+    scopes_str = ENTRAID_CFG["scopes"]
+    scopes = tuple(scope.strip() for scope in scopes_str.split(","))
+
+    return create_from_default_azure_credential(
+        scopes=scopes, token_manager_config=token_manager_config
+    )
diff --git a/src/main.py b/src/main.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@ dependencies = [`
`26`	`26`	`"dotenv>=0.9.9",`
`27`	`27`	`"numpy>=2.2.4",`
`28`	`28`	`"click>=8.0.0",`
	`29`	`+ "redis-entraid>=1.0.0",`
`29`	`30`	`]`
`30`	`31`
`31`	`32`	`[project.scripts]`