Merge pull request #524 from microsoft/psl-sfi

Prajwal-Microsoft · web-flow · commit b8f300f81dba · 2025-07-31T08:55:42.000+05:30
feat: replacing DefaultAzureCredential with ManagedIdentityCredential
diff --git a/infra/deploy_app_service.bicep b/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
     }
diff --git a/infra/main.json b/infra/main.json
@@ -5,7 +5,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.36.177.2456",
-      "templateHash": "7421251462011771854"
+      "templateHash": "12338844521177910469"
     }
   },
   "parameters": {
@@ -1848,7 +1848,7 @@
             "_generator": {
               "name": "bicep",
               "version": "0.36.177.2456",
-              "templateHash": "16850060889438240970"
+              "templateHash": "7711665754275271950"
             }
           },
           "parameters": {
@@ -2308,6 +2308,10 @@
                     {
                       "name": "UWSGI_THREADS",
                       "value": "2"
+                    },
+                    {
+                      "name": "APP_ENV",
+                      "value": "Prod"
                     }
                   ],
                   "linuxFxVersion": "[variables('imageName')]"
diff --git a/infra/scripts/index_scripts/01_create_search_index.py b/infra/scripts/index_scripts/01_create_search_index.py
@@ -1,4 +1,4 @@
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 from azure.search.documents.indexes import SearchIndexClient
 from azure.search.documents.indexes.models import (
@@ -27,11 +27,11 @@ def get_secrets_from_kv(secret_name: str) -> str:
     Retrieves a secret value from Azure Key Vault.
     Args:
         secret_name (str): Name of the secret.
-        credential (DefaultAzureCredential): Credential with access to Key Vault.
+        credential (AzureCliCredential): Credential with access to Key Vault.
     Returns:
         str: The secret value.
     """
-    kv_credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    kv_credential = AzureCliCredential()
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/",
         credential=kv_credential
@@ -43,7 +43,7 @@ def create_search_index():
     """Create an Azure Search index."""
 
     # Shared credential
-    credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    credential = AzureCliCredential()
 
     # Retrieve secrets from Key Vault
     search_endpoint = get_secrets_from_kv("AZURE-SEARCH-ENDPOINT")
diff --git a/infra/scripts/index_scripts/02_process_data.py b/infra/scripts/index_scripts/02_process_data.py
@@ -7,7 +7,7 @@
 from azure.search.documents import SearchClient
 from azure.storage.filedatalake import DataLakeServiceClient
 from azure.search.documents.indexes import SearchIndexClient
-from azure.identity import (DefaultAzureCredential, get_bearer_token_provider)
+from azure.identity import (AzureCliCredential, get_bearer_token_provider)
 
 
 key_vault_name = 'kv_to-be-replaced'
@@ -22,11 +22,11 @@ def get_secrets_from_kv(secret_name: str) -> str:
     Retrieves a secret value from Azure Key Vault.
     Args:
         secret_name (str): Name of the secret.
-        credential (DefaultAzureCredential): Credential with access to Key Vault.
+        credential (AzureCliCredential): Credential with access to Key Vault.
     Returns:
         str: The secret value.
     """
-    kv_credential = DefaultAzureCredential(managed_identity_client_id=managed_identity_client_id)
+    kv_credential = AzureCliCredential()
     secret_client = SecretClient(
         vault_url=f"https://{key_vault_name}.vault.azure.net/",
         credential=kv_credential
@@ -44,7 +44,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)
+credential = AzureCliCredential()
 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
diff --git a/scripts/chunk_documents.py b/scripts/chunk_documents.py
@@ -5,7 +5,7 @@
 
 from azure.ai.formrecognizer import DocumentAnalysisClient
 from azure.core.credentials import AzureKeyCredential
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 from data_utils import chunk_directory
 
@@ -57,8 +57,7 @@ def get_document_intelligence_client(config, secret_client):
     with open(args.config_file) as f:
         config = json.load(f)
 
-    credential = DefaultAzureCredential()
-
+    credential = AzureCliCredential()
     if type(config) is not list:
         config = [config]
 
diff --git a/scripts/embed_documents.py b/scripts/embed_documents.py
@@ -2,7 +2,7 @@
 import json
 from asyncio import sleep
 
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureCliCredential
 from azure.keyvault.secrets import SecretClient
 from data_utils import get_embedding
 
@@ -19,8 +19,7 @@
     with open(args.config_file) as f:
         config = json.load(f)
 
-    credential = DefaultAzureCredential()
-
+    credential = AzureCliCredential()
     if type(config) is not list:
         config = [config]
 
diff --git a/src/app.py b/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
 
diff --git a/src/backend/api/agent/browse_agent_factory.py b/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
         )
 
diff --git a/src/backend/api/agent/section_agent_factory.py b/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
         )
 
diff --git a/src/backend/api/agent/template_agent_factory.py b/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
         )
 
diff --git a/src/backend/helpers/azure_credential_utils.py b/src/backend/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)

Original file line number	Diff line number	Diff line change
`@@ -298,6 +298,10 @@ resource Website 'Microsoft.Web/sites@2020-06-01' = {`
`298`	`298`	`name: 'UWSGI_THREADS'`
`299`	`299`	`value: '2'`
`300`	`300`	`}`
	`301`	`+ {`
	`302`	`+ name: 'APP_ENV'`
	`303`	`+ value: 'Prod'`
	`304`	`+ }`
`301`	`305`	`]`
`302`	`306`	`linuxFxVersion: imageName`
`303`	`307`	`}`