fix: Replace DefaultAzureCredential with ManagedIdentityCredential for production-safe authentication

.env.sample

@@ -60,6 +60,8 @@ AZURE_SPEECH_SERVICE_REGION=
 AZURE_AUTH_TYPE=keys
 USE_KEY_VAULT=true
 AZURE_KEY_VAULT_ENDPOINT=
+# Application environment (e.g., dev, prod)
+APP_ENV="dev"
 # Chat conversation type to decide between custom or byod (bring your own data) conversation type
 CONVERSATION_FLOW=
 # Chat History CosmosDB Integration Settings

code/backend/batch/utilities/chat_history/database_factory.py

@@ -2,7 +2,7 @@
 from ..helpers.env_helper import EnvHelper
 from .cosmosdb import CosmosConversationClient
 from .postgresdbservice import PostgresConversationClient
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from ..helpers.config.database_type import DatabaseType
 
 
@@ -25,7 +25,7 @@ def get_conversation_client():
                 f"https://{env_helper.AZURE_COSMOSDB_ACCOUNT}.documents.azure.com:443/"
             )
             credential = (
-                DefaultAzureCredential()
+                get_azure_credential()
                 if not env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
                 else env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
             )

code/backend/batch/utilities/chat_history/postgresdbservice.py

@@ -1,7 +1,7 @@
 import logging
 import asyncpg
 from datetime import datetime, timezone
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 
 from .database_client_base import DatabaseClientBase
 
@@ -21,7 +21,7 @@ def __init__(
 
     async def connect(self):
         try:
-            credential = DefaultAzureCredential()
+            credential = get_azure_credential()
             token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             ).token

code/backend/batch/utilities/helpers/azure_blob_storage_client.py

@@ -12,7 +12,7 @@
 from azure.storage.queue import QueueClient, BinaryBase64EncodePolicy
 import chardet
 from .env_helper import EnvHelper
-from azure.identity import DefaultAzureCredential
+from .azure_credential_utils import get_azure_credential
 
 
 def connection_string(account_name: str, account_key: str):
@@ -25,7 +25,7 @@ def create_queue_client():
         return QueueClient(
             account_url=f"https://{env_helper.AZURE_BLOB_ACCOUNT_NAME}.queue.core.windows.net/",
             queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME,
-            credential=DefaultAzureCredential(),
+            credential=get_azure_credential(),
             message_encode_policy=BinaryBase64EncodePolicy(),
         )
 
@@ -56,7 +56,7 @@ def __init__(
         if self.auth_type == "rbac":
             self.account_key = None
             self.blob_service_client = BlobServiceClient(
-                account_url=self.endpoint, credential=DefaultAzureCredential()
+                account_url=self.endpoint, credential=get_azure_credential()
             )
             self.user_delegation_key = self.request_user_delegation_key(
                 blob_service_client=self.blob_service_client

code/backend/batch/utilities/helpers/azure_computer_vision_client.py

@@ -1,6 +1,7 @@
 import logging
 from urllib.parse import urljoin
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from azure.identity import get_bearer_token_provider
+from .azure_credential_utils import get_azure_credential
 
 import requests
 from requests import Response
@@ -56,7 +57,7 @@ def __make_request(self, path: str, body) -> Response:
                 headers["Ocp-Apim-Subscription-Key"] = self.key
             else:
                 token_provider = get_bearer_token_provider(
-                    DefaultAzureCredential(), self.__TOKEN_SCOPE
+                    get_azure_credential(), self.__TOKEN_SCOPE
                 )
                 headers["Authorization"] = "Bearer " + token_provider()
 

code/backend/batch/utilities/helpers/azure_credential_utils.py

@@ -0,0 +1,48 @@
+import os
+from azure.identity import ManagedIdentityCredential, DefaultAzureCredential
+from azure.identity.aio import (
+    ManagedIdentityCredential as AioManagedIdentityCredential,
+    DefaultAzureCredential as AioDefaultAzureCredential,
+)
+
+
+async def get_azure_credential_async(client_id=None):
+    """
+    Returns an Azure credential asynchronously based on the application environment.
+
+    If the environment is 'dev', it uses AioDefaultAzureCredential.
+    Otherwise, it uses AioManagedIdentityCredential.
+
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+
+    Returns:
+        Credential object: Either AioDefaultAzureCredential or AioManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == "dev":
+        return (
+            AioDefaultAzureCredential()
+        )  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return AioManagedIdentityCredential(client_id=client_id)
+
+
+def get_azure_credential(client_id=None):
+    """
+    Returns an Azure credential based on the application environment.
+
+    If the environment is 'dev', it uses DefaultAzureCredential.
+    Otherwise, it uses ManagedIdentityCredential.
+
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+
+    Returns:
+        Credential object: Either DefaultAzureCredential or ManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == "dev":
+        return (
+            DefaultAzureCredential()
+        )  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return ManagedIdentityCredential(client_id=client_id)

code/backend/batch/utilities/helpers/azure_form_recognizer_helper.py

@@ -1,7 +1,7 @@
 import logging
 from azure.core.credentials import AzureKeyCredential
 from azure.ai.formrecognizer import DocumentAnalysisClient
-from azure.identity import DefaultAzureCredential
+from .azure_credential_utils import get_azure_credential
 import html
 import traceback
 from .env_helper import EnvHelper
@@ -19,7 +19,7 @@ def __init__(self) -> None:
         if env_helper.AZURE_AUTH_TYPE == "rbac":
             self.document_analysis_client = DocumentAnalysisClient(
                 endpoint=self.AZURE_FORM_RECOGNIZER_ENDPOINT,
-                credential=DefaultAzureCredential(),
+                credential=get_azure_credential(),
                 headers={
                     "x-ms-useragent": "chat-with-your-data-solution-accelerator/1.0.0"
                 },

code/backend/batch/utilities/helpers/azure_postgres_helper.py

@@ -1,7 +1,7 @@
 import logging
 import psycopg2
 from psycopg2.extras import execute_values, RealDictCursor
-from azure.identity import DefaultAzureCredential
+from .azure_credential_utils import get_azure_credential
 from .llm_helper import LLMHelper
 from .env_helper import EnvHelper
 
@@ -24,7 +24,7 @@ def _create_search_client(self):
             dbname = self.env_helper.POSTGRESQL_DATABASE
 
             # Acquire the access token
-            credential = DefaultAzureCredential()
+            credential = get_azure_credential()
             access_token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             )

code/backend/batch/utilities/helpers/azure_search_helper.py

@@ -2,7 +2,7 @@
 from typing import Union
 from langchain_community.vectorstores import AzureSearch
 from azure.core.credentials import AzureKeyCredential
-from azure.identity import DefaultAzureCredential
+from .azure_credential_utils import get_azure_credential
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
 from azure.search.documents.indexes.models import (
@@ -49,10 +49,10 @@ def _search_credential(self):
         if self.env_helper.is_auth_type_keys():
             return AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
         else:
-            return DefaultAzureCredential()
+            return get_azure_credential()
 
     def _create_search_client(
-        self, search_credential: Union[AzureKeyCredential, DefaultAzureCredential]
+        self, search_credential: Union[AzureKeyCredential, get_azure_credential]
     ) -> SearchClient:
         return SearchClient(
             endpoint=self.env_helper.AZURE_SEARCH_SERVICE,
@@ -61,7 +61,7 @@ def _create_search_client(
         )
 
     def _create_search_index_client(
-        self, search_credential: Union[AzureKeyCredential, DefaultAzureCredential]
+        self, search_credential: Union[AzureKeyCredential, get_azure_credential]
     ):
         return SearchIndexClient(
             endpoint=self.env_helper.AZURE_SEARCH_SERVICE, credential=search_credential
@@ -285,7 +285,7 @@ def get_conversation_logger(self):
         ]
 
         if self.env_helper.AZURE_AUTH_TYPE == "rbac":
-            credential = DefaultAzureCredential()
+            credential = get_azure_credential()
             return AzureSearch(
                 azure_search_endpoint=self.env_helper.AZURE_SEARCH_SERVICE,
                 azure_search_key=None,  # Remove API key

code/backend/batch/utilities/helpers/env_helper.py

@@ -3,9 +3,9 @@
 import logging
 import threading
 from dotenv import load_dotenv
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from azure.identity import get_bearer_token_provider
+from .azure_credential_utils import get_azure_credential
 from azure.keyvault.secrets import SecretClient
-
 from ..orchestrator.orchestration_strategy import OrchestrationStrategy
 from ..helpers.config.conversation_flow import ConversationFlow
 from ..helpers.config.database_type import DatabaseType
@@ -216,7 +216,7 @@ def __load_config(self, **kwargs) -> None:
         )
 
         self.AZURE_TOKEN_PROVIDER = get_bearer_token_provider(
-            DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            get_azure_credential(), "https://cognitiveservices.azure.com/.default"
         )
         self.ADVANCED_IMAGE_PROCESSING_MAX_IMAGES = self.get_env_var_int(
             "ADVANCED_IMAGE_PROCESSING_MAX_IMAGES", 1
@@ -362,8 +362,8 @@ def __load_config(self, **kwargs) -> None:
         self.OPEN_AI_FUNCTIONS_SYSTEM_PROMPT = os.getenv(
             "OPEN_AI_FUNCTIONS_SYSTEM_PROMPT", ""
         )
-        self.SEMENTIC_KERNEL_SYSTEM_PROMPT = os.getenv(
-            "SEMENTIC_KERNEL_SYSTEM_PROMPT", ""
+        self.SEMANTIC_KERNEL_SYSTEM_PROMPT = os.getenv(
+            "SEMANTIC_KERNEL_SYSTEM_PROMPT", ""
         )
 
         self.ENFORCE_AUTH = self.get_env_var_bool("ENFORCE_AUTH", "True")
@@ -429,7 +429,7 @@ def __init__(self) -> None:
         self.secret_client = None
         if self.USE_KEY_VAULT:
             self.secret_client = SecretClient(
-                os.environ.get("AZURE_KEY_VAULT_ENDPOINT"), DefaultAzureCredential()
+                os.environ.get("AZURE_KEY_VAULT_ENDPOINT"), get_azure_credential()
             )
 
     def get_secret(self, secret_name: str) -> str:

code/backend/batch/utilities/helpers/llm_helper.py

@@ -8,7 +8,7 @@
     AzureChatPromptExecutionSettings,
 )
 from azure.ai.ml import MLClient
-from azure.identity import DefaultAzureCredential
+from .azure_credential_utils import get_azure_credential
 from .env_helper import EnvHelper
 
 logger = logging.getLogger(__name__)
@@ -166,7 +166,7 @@ def get_sk_service_settings(self, service: AzureChatCompletion):
     def get_ml_client(self):
         if not hasattr(self, "_ml_client"):
             self._ml_client = MLClient(
-                DefaultAzureCredential(),
+                get_azure_credential(),
                 self.env_helper.AZURE_SUBSCRIPTION_ID,
                 self.env_helper.AZURE_RESOURCE_GROUP,
                 self.env_helper.AZURE_ML_WORKSPACE_NAME,

code/backend/batch/utilities/integrated_vectorization/azure_search_datasource.py

@@ -7,7 +7,7 @@
 )
 from azure.search.documents.indexes import SearchIndexerClient
 from ..helpers.env_helper import EnvHelper
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from azure.core.credentials import AzureKeyCredential
 
 
@@ -19,7 +19,7 @@ def __init__(self, env_helper: EnvHelper):
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else DefaultAzureCredential()
+                else get_azure_credential()
             ),
         )
 

code/backend/batch/utilities/integrated_vectorization/azure_search_index.py

@@ -21,7 +21,7 @@
     SearchIndex,
 )
 from ..helpers.env_helper import EnvHelper
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from azure.core.credentials import AzureKeyCredential
 from ..helpers.llm_helper import LLMHelper
 
@@ -39,7 +39,7 @@ def __init__(self, env_helper: EnvHelper, llm_helper: LLMHelper):
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else DefaultAzureCredential()
+                else get_azure_credential()
             ),
         )
 

code/backend/batch/utilities/integrated_vectorization/azure_search_indexer.py

@@ -2,7 +2,7 @@
 from azure.search.documents.indexes.models import SearchIndexer, FieldMapping
 from azure.search.documents.indexes import SearchIndexerClient
 from ..helpers.env_helper import EnvHelper
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from azure.core.credentials import AzureKeyCredential
 
 logger = logging.getLogger(__name__)
@@ -16,7 +16,7 @@ def __init__(self, env_helper: EnvHelper):
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else DefaultAzureCredential()
+                else get_azure_credential()
             ),
         )
 

code/backend/batch/utilities/integrated_vectorization/azure_search_skillset.py

@@ -15,7 +15,7 @@
 from azure.search.documents.indexes import SearchIndexerClient
 from ..helpers.config.config_helper import IntegratedVectorizationConfig
 from ..helpers.env_helper import EnvHelper
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from azure.core.credentials import AzureKeyCredential
 
 logger = logging.getLogger(__name__)
@@ -33,7 +33,7 @@ def __init__(
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else DefaultAzureCredential()
+                else get_azure_credential()
             ),
         )
         self.integrated_vectorization_config = integrated_vectorization_config

code/backend/batch/utilities/orchestrator/semantic_kernel.py

@@ -41,7 +41,7 @@ async def orchestrate(
             if response := self.call_content_safety_input(user_message):
                 return response
 
-        system_message = self.env_helper.SEMENTIC_KERNEL_SYSTEM_PROMPT
+        system_message = self.env_helper.SEMANTIC_KERNEL_SYSTEM_PROMPT
         if not system_message:
             system_message = """You help employees to navigate only private information sources.
 You must prioritize the function call over your general knowledge for any question by calling the search_documents function.

code/backend/batch/utilities/search/integrated_vectorization_search_handler.py

@@ -5,7 +5,7 @@
 from azure.search.documents.indexes import SearchIndexClient
 from azure.search.documents.models import VectorizableTextQuery
 from azure.core.credentials import AzureKeyCredential
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from ..common.source_document import SourceDocument
 import re
 
@@ -21,7 +21,7 @@ def create_search_client(self):
                 credential=(
                     AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                     if self.env_helper.is_auth_type_keys()
-                    else DefaultAzureCredential()
+                    else get_azure_credential()
                 ),
             )
 
@@ -170,7 +170,7 @@ def _check_index_exists(self) -> bool:
             credential=(
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else DefaultAzureCredential()
+                else get_azure_credential()
             ),
         )
 

code/backend/batch/utilities/tools/content_safety_checker.py

@@ -1,7 +1,7 @@
 import logging
 from azure.ai.contentsafety import ContentSafetyClient
 from azure.core.credentials import AzureKeyCredential
-from azure.identity import DefaultAzureCredential
+from ..helpers.azure_credential_utils import get_azure_credential
 from azure.core.exceptions import HttpResponseError
 from azure.ai.contentsafety.models import AnalyzeTextOptions
 from ..helpers.env_helper import EnvHelper
@@ -19,7 +19,7 @@ def __init__(self):
             logger.info("Initializing ContentSafetyClient with RBAC authentication.")
             self.content_safety_client = ContentSafetyClient(
                 env_helper.AZURE_CONTENT_SAFETY_ENDPOINT,
-                DefaultAzureCredential(),
+                get_azure_credential(),
             )
         else:
             logger.info(