replacing DefaultAzureCredential with ManagedIdentityCredential

Author: Vamshi-Microsoft

infra/deploy_app_service.bicep

@@ -298,6 +298,10 @@ resource Website 'Microsoft.Web/sites@2020-06-01' = {
           name: 'UWSGI_THREADS'
           value: '2'
         }
+        {
+          name: 'APP_ENV'
+          value: 'Prod'
+        }
       ]
       linuxFxVersion: imageName
     }

infra/scripts/index_scripts/01_create_search_index.py

@@ -32,6 +32,7 @@ def get_secrets_from_kv(secret_name: str) -> str:
         str: The secret value.
     """
     kv_credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/",
         credential=kv_credential
@@ -44,6 +45,7 @@ def create_search_index():
 
     # Shared credential
     credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
 
     # Retrieve secrets from Key Vault
     search_endpoint = get_secrets_from_kv("AZURE-SEARCH-ENDPOINT")

infra/scripts/index_scripts/02_process_data.py

@@ -27,6 +27,7 @@ def get_secrets_from_kv(secret_name: str) -> str:
         str: The secret value.
     """
     kv_credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/",
         credential=kv_credential
@@ -45,6 +46,7 @@ def get_secrets_from_kv(secret_name: str) -> str:
 # Azure Data Lake settings
 account_url = f"https://{account_name}.dfs.core.windows.net"
 credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+# CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
 service_client = DataLakeServiceClient(account_url, credential=credential, api_version='2023-01-03')
 file_system_client = service_client.get_file_system_client(file_system_client_name)
 directory_name = directory

scripts/chunk_documents.py

@@ -58,6 +58,7 @@ def get_document_intelligence_client(config, secret_client):
         config = json.load(f)
 
     credential = DefaultAzureCredential()
+    # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
 
     if type(config) is not list:
         config = [config]

scripts/embed_documents.py

@@ -20,6 +20,7 @@
         config = json.load(f)
 
     credential = DefaultAzureCredential()
+    # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in local environment.
 
     if type(config) is not list:
         config = [config]

src/app.py

@@ -8,8 +8,9 @@
 import asyncio
 from typing import Dict, Any, AsyncGenerator
 
-from azure.identity.aio import DefaultAzureCredential
-from azure.identity import DefaultAzureCredential as DefaultAzureCredentialSync
+
+from backend.helpers.azure_credential_utils import get_azure_credential
+from backend.helpers.azure_credential_utils import get_azure_credential_async
 from quart import (Blueprint, Quart, jsonify, make_response, render_template,
                    request, send_from_directory)
 
@@ -169,7 +170,7 @@ async def init_ai_foundry_client():
 
         ai_project_client = AIProjectClient(
             endpoint=app_settings.azure_ai.agent_endpoint,
-            credential=DefaultAzureCredential()
+            credential=get_azure_credential()
         )
         track_event_if_configured("AIFoundryAgentEndpointUsed", {
             "endpoint": app_settings.azure_ai.agent_endpoint
@@ -193,7 +194,7 @@ def init_cosmosdb_client():
             )
 
             if not app_settings.chat_history.account_key:
-                credential = DefaultAzureCredential()
+                credential = get_azure_credential()
             else:
                 credential = app_settings.chat_history.account_key
 
@@ -1163,7 +1164,7 @@ async def fetch_azure_search_content():
             return jsonify({"error": "URL and title are required"}), 400
 
         # Get Azure AD token
-        credential = DefaultAzureCredentialSync()
+        credential = await get_azure_credential_async()
         token = credential.get_token("https://search.azure.com/.default")
         access_token = token.token
 

src/backend/api/agent/browse_agent_factory.py

@@ -1,6 +1,6 @@
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.agents.models import AzureAISearchTool, AzureAISearchQueryType
-from azure.identity.aio import DefaultAzureCredential
+from backend.helpers.azure_credential_utils import get_azure_credential
 from backend.settings import app_settings
 from event_utils import track_event_if_configured
 
@@ -20,7 +20,7 @@ async def create_or_get_agent(cls):
         """
         project_client = AIProjectClient(
             endpoint=app_settings.azure_ai.agent_endpoint,
-            credential=DefaultAzureCredential(exclude_interactive_browser_credential=False),
+            credential=get_azure_credential(),
             api_version=app_settings.azure_ai.agent_api_version
         )
 

src/backend/api/agent/section_agent_factory.py

@@ -1,6 +1,6 @@
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.agents.models import AzureAISearchTool, AzureAISearchQueryType
-from azure.identity.aio import DefaultAzureCredential
+from backend.helpers.azure_credential_utils import get_azure_credential
 from backend.settings import app_settings
 from event_utils import track_event_if_configured
 
@@ -19,7 +19,7 @@ async def create_or_get_agent(cls):
         """
         project_client = AIProjectClient(
             endpoint=app_settings.azure_ai.agent_endpoint,
-            credential=DefaultAzureCredential(exclude_interactive_browser_credential=False),
+            credential=get_azure_credential(),
             api_version=app_settings.azure_ai.agent_api_version
         )
 

src/backend/api/agent/template_agent_factory.py

@@ -1,6 +1,6 @@
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.agents.models import AzureAISearchTool, AzureAISearchQueryType
-from azure.identity.aio import DefaultAzureCredential
+from backend.helpers.azure_credential_utils import get_azure_credential
 from backend.settings import app_settings
 from event_utils import track_event_if_configured
 
@@ -20,7 +20,7 @@ async def create_or_get_agent(cls):
         """
         project_client = AIProjectClient(
             endpoint=app_settings.azure_ai.agent_endpoint,
-            credential=DefaultAzureCredential(exclude_interactive_browser_credential=False),
+            credential=get_azure_credential(),
             api_version=app_settings.azure_ai.agent_api_version
         )
 

src/helpers/azure_credential_utils.py

@@ -0,0 +1,41 @@
+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)