feat: WAF implementation for CWYD

.github/workflows/group_dependabot_security_updates.yml

@@ -65,9 +65,10 @@ jobs:
           git config --global user.email "[email protected]"
 
       - name: Install required tools
-        uses: awalsh128/[email protected]
-        with:
-          packages: "jq gh"
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y jq gh
+        shell: bash
 
       - name: Enable strict error handling
         shell: bash

azure.yaml

@@ -2,7 +2,7 @@
 
 name: chat-with-your-data-solution-accelerator
 metadata:
-  template: [email protected]
+   template: [email protected]
 hooks:
   postprovision:
     # run: ./infra/prompt-flow/create-prompt-flow.sh

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

@@ -25,7 +25,7 @@ def get_conversation_client():
                 f"https://{env_helper.AZURE_COSMOSDB_ACCOUNT}.documents.azure.com:443/"
             )
             credential = (
-                get_azure_credential()
+                get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID)
                 if not env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
                 else env_helper.AZURE_COSMOSDB_ACCOUNT_KEY
             )

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

@@ -2,6 +2,7 @@
 import asyncpg
 from datetime import datetime, timezone
 from ..helpers.azure_credential_utils import get_azure_credential
+from ..helpers.env_helper import EnvHelper
 
 from .database_client_base import DatabaseClientBase
 
@@ -13,6 +14,7 @@ class PostgresConversationClient(DatabaseClientBase):
     def __init__(
         self, user: str, host: str, database: str, enable_message_feedback: bool = False
     ):
+        self.env_helper = EnvHelper()
         self.user = user
         self.host = host
         self.database = database
@@ -21,7 +23,7 @@ def __init__(
 
     async def connect(self):
         try:
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             ).token
@@ -31,7 +33,7 @@ async def connect(self):
                 database=self.database,
                 password=token,
                 port=5432,
-                ssl="require",
+                ssl=True,
             )
         except Exception as e:
             logger.error("Failed to connect to PostgreSQL: %s", e)

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

@@ -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=get_azure_credential(),
+            credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID),
             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=get_azure_credential()
+                account_url=self.endpoint, credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID)
             )
             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

@@ -27,6 +27,7 @@ def __init__(self, env_helper: EnvHelper) -> None:
         self.model_version = (
             env_helper.AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION
         )
+        self.managed_identity_client_id = env_helper.MANAGED_IDENTITY_CLIENT_ID
 
     def vectorize_image(self, image_url: str) -> list[float]:
         logger.info(f"Making call to computer vision to vectorize image: {image_url}")
@@ -57,7 +58,7 @@ def __make_request(self, path: str, body) -> Response:
                 headers["Ocp-Apim-Subscription-Key"] = self.key
             else:
                 token_provider = get_bearer_token_provider(
-                    get_azure_credential(), self.__TOKEN_SCOPE
+                    get_azure_credential(self.managed_identity_client_id), self.__TOKEN_SCOPE
                 )
                 headers["Authorization"] = "Bearer " + token_provider()
 

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

@@ -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=get_azure_credential(),
+                credential=get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID),
                 headers={
                     "x-ms-useragent": "chat-with-your-data-solution-accelerator/1.0.0"
                 },

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

@@ -24,14 +24,14 @@ def _create_search_client(self):
             dbname = self.env_helper.POSTGRESQL_DATABASE
 
             # Acquire the access token
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             access_token = credential.get_token(
                 "https://ossrdbms-aad.database.windows.net/.default"
             )
 
             # Use the token in the connection string
             conn_string = (
-                f"host={host} user={user} dbname={dbname} password={access_token.token}"
+                f"host={host} user={user} dbname={dbname} password={access_token.token} sslmode=require"
             )
             self.conn = psycopg2.connect(conn_string)
             logger.info("Connected to Azure PostgreSQL successfully.")

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

@@ -49,7 +49,7 @@ def _search_credential(self):
         if self.env_helper.is_auth_type_keys():
             return AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
         else:
-            return get_azure_credential()
+            return get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
 
     def _create_search_client(
         self, search_credential: Union[AzureKeyCredential, get_azure_credential]
@@ -285,7 +285,7 @@ def get_conversation_logger(self):
         ]
 
         if self.env_helper.AZURE_AUTH_TYPE == "rbac":
-            credential = get_azure_credential()
+            credential = get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             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

@@ -35,10 +35,13 @@ def __load_config(self, **kwargs) -> None:
         self.secretHelper = SecretHelper()
 
         self.LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
+        self.APP_ENV = os.getenv("APP_ENV", "Prod").lower()
 
         # Azure
         self.AZURE_SUBSCRIPTION_ID = os.getenv("AZURE_SUBSCRIPTION_ID", "")
         self.AZURE_RESOURCE_GROUP = os.getenv("AZURE_RESOURCE_GROUP", "")
+        self.MANAGED_IDENTITY_CLIENT_ID = os.getenv("MANAGED_IDENTITY_CLIENT_ID", "")
+        self.MANAGED_IDENTITY_RESOURCE_ID = os.getenv("MANAGED_IDENTITY_RESOURCE_ID", "")
 
         # Azure Search
         self.AZURE_SEARCH_SERVICE = os.getenv("AZURE_SEARCH_SERVICE", "")
@@ -217,7 +220,7 @@ def __load_config(self, **kwargs) -> None:
         )
 
         self.AZURE_TOKEN_PROVIDER = get_bearer_token_provider(
-            get_azure_credential(), "https://cognitiveservices.azure.com/.default"
+            get_azure_credential(self.MANAGED_IDENTITY_CLIENT_ID), "https://cognitiveservices.azure.com/.default"
         )
         self.ADVANCED_IMAGE_PROCESSING_MAX_IMAGES = self.get_env_var_int(
             "ADVANCED_IMAGE_PROCESSING_MAX_IMAGES", 1
@@ -234,7 +237,7 @@ def __load_config(self, **kwargs) -> None:
         self.AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION = os.getenv(
             "AZURE_COMPUTER_VISION_VECTORIZE_IMAGE_MODEL_VERSION", "2023-04-15"
         )
-        self.FUNCTION_KEY = os.getenv("FUNCTION_KEY", "")
+        self.FUNCTION_KEY = self.secretHelper.get_secret("FUNCTION_KEY")
 
         # Initialize Azure keys based on authentication type and environment settings.
         # When AZURE_AUTH_TYPE is "rbac", azure keys are None or an empty string.
@@ -243,7 +246,6 @@ def __load_config(self, **kwargs) -> None:
             self.AZURE_OPENAI_API_KEY = ""
             self.AZURE_SPEECH_KEY = None
             self.AZURE_COMPUTER_VISION_KEY = None
-            self.FUNCTION_KEY = self.secretHelper.get_secret("FUNCTION_KEY")
         else:
             self.AZURE_SEARCH_KEY = self.secretHelper.get_secret("AZURE_SEARCH_KEY")
             self.AZURE_OPENAI_API_KEY = self.secretHelper.get_secret(
@@ -429,8 +431,11 @@ def __init__(self) -> None:
         self.USE_KEY_VAULT = os.getenv("USE_KEY_VAULT", "").lower() == "true"
         self.secret_client = None
         if self.USE_KEY_VAULT:
+            vault_endpoint = os.environ.get("AZURE_KEY_VAULT_ENDPOINT")
+            if not vault_endpoint:
+                raise ValueError("AZURE_KEY_VAULT_ENDPOINT environment variable is required when USE_KEY_VAULT is true")
             self.secret_client = SecretClient(
-                os.environ.get("AZURE_KEY_VAULT_ENDPOINT"), get_azure_credential()
+                vault_endpoint, get_azure_credential(client_id=os.getenv("MANAGED_IDENTITY_CLIENT_ID", None))
             )
 
     def get_secret(self, secret_name: str) -> str:

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

@@ -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(
-                get_azure_credential(),
+                get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID),
                 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

@@ -1,6 +1,7 @@
 from azure.search.documents.indexes.models import (
     SearchIndexerDataContainer,
     SearchIndexerDataSourceConnection,
+    SearchIndexerDataUserAssignedIdentity,
 )
 from azure.search.documents.indexes._generated.models import (
     NativeBlobSoftDeleteDeletionDetectionPolicy,
@@ -19,7 +20,7 @@ def __init__(self, env_helper: EnvHelper):
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else get_azure_credential()
+                else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             ),
         )
 
@@ -35,6 +36,13 @@ def create_or_update_datasource(self):
             connection_string=connection_string,
             container=container,
             data_deletion_detection_policy=NativeBlobSoftDeleteDeletionDetectionPolicy(),
+            identity=(
+                None
+                if getattr(self.env_helper, "APP_ENV", "").lower() == "dev"
+                else SearchIndexerDataUserAssignedIdentity(
+                    user_assigned_identity=self.env_helper.MANAGED_IDENTITY_RESOURCE_ID
+                )
+            ),
         )
         self.indexer_client.create_or_update_data_source_connection(
             data_source_connection

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

@@ -19,6 +19,7 @@
     SemanticPrioritizedFields,
     SemanticField,
     SearchIndex,
+    SearchIndexerDataUserAssignedIdentity,
 )
 from ..helpers.env_helper import EnvHelper
 from ..helpers.azure_credential_utils import get_azure_credential
@@ -39,7 +40,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 get_azure_credential()
+                else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             ),
         )
 
@@ -144,6 +145,13 @@ def get_vector_search_config(self):
             azure_open_ai_parameters = AzureOpenAIParameters(
                 resource_uri=self.env_helper.AZURE_OPENAI_ENDPOINT,
                 deployment_id=self.env_helper.AZURE_OPENAI_EMBEDDING_MODEL,
+                auth_identity=(
+                    None
+                    if getattr(self.env_helper, "APP_ENV", "").lower() == "dev"
+                    else SearchIndexerDataUserAssignedIdentity(
+                        user_assigned_identity=self.env_helper.MANAGED_IDENTITY_RESOURCE_ID
+                    )
+                ),
             )
 
         return VectorSearch(

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

@@ -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 get_azure_credential()
+                else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             ),
         )
 

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

@@ -11,6 +11,7 @@
     SearchIndexerIndexProjectionsParameters,
     IndexProjectionMode,
     SearchIndexerSkillset,
+    SearchIndexerDataUserAssignedIdentity,
 )
 from azure.search.documents.indexes import SearchIndexerClient
 from ..helpers.config.config_helper import IntegratedVectorizationConfig
@@ -33,7 +34,7 @@ def __init__(
             (
                 AzureKeyCredential(self.env_helper.AZURE_SEARCH_KEY)
                 if self.env_helper.is_auth_type_keys()
-                else get_azure_credential()
+                else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             ),
         )
         self.integrated_vectorization_config = integrated_vectorization_config
@@ -95,6 +96,13 @@ def create_skillset(self):
                 if self.env_helper.is_auth_type_keys()
                 else None
             ),
+            auth_identity=(
+                None
+                if getattr(self.env_helper, "APP_ENV", "").lower() == "dev"
+                else SearchIndexerDataUserAssignedIdentity(
+                    user_assigned_identity=self.env_helper.MANAGED_IDENTITY_RESOURCE_ID
+                )
+            ),
             inputs=[
                 InputFieldMappingEntry(name="text", source="/document/pages/*"),
             ],

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

@@ -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 get_azure_credential()
+                    else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
                 ),
             )
 
@@ -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 get_azure_credential()
+                else get_azure_credential(self.env_helper.MANAGED_IDENTITY_CLIENT_ID)
             ),
         )
 

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

@@ -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,
-                get_azure_credential(),
+                get_azure_credential(env_helper.MANAGED_IDENTITY_CLIENT_ID),
             )
         else:
             logger.info(

code/backend/pages/01_Ingest_Data.py

@@ -55,7 +55,10 @@ def reprocess_all():
 
 def add_urls():
     urls = st.session_state["urls"].split("\n")
-    add_url_embeddings(urls)
+    result = add_url_embeddings(urls)
+    # If URLs are valid and processed, clear the textarea
+    if result:
+        st.session_state["urls"] = ""
 
 
 def sanitize_metadata_value(value):
@@ -67,7 +70,7 @@ def add_url_embeddings(urls: list[str]):
     has_valid_url = bool(list(filter(str.strip, urls)))
     if not has_valid_url:
         st.error("Please enter at least one valid URL.")
-        return
+        return False
 
     params = {}
     if env_helper.FUNCTION_KEY is not None:
@@ -80,9 +83,11 @@ def add_url_embeddings(urls: list[str]):
         )
         r = requests.post(url=backend_url, params=params, json=body)
         if not r.ok:
-            raise ValueError(f"Error {r.status_code}: {r.text}")
+            st.error(f"Error {r.status_code}: {r.text}")
+            return False
         else:
             st.success(f"Embeddings added successfully for {url}")
+    return True
 
 
 try: