fixing errors and permissions

Author: cyphronix

aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py

@@ -72,15 +72,16 @@ def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None:  # ty
         encryption_config = domain.get("DomainStatus", {}).get("EncryptionAtRestOptions", {})
         if not encryption_config.get("Enabled", False):
             return f"{kb_name} (OpenSearch domain encryption not enabled)"
-        if not encryption_config.get("KmsKeyId"):
+        kms_key_id = encryption_config.get("KmsKeyId", "")
+        if not kms_key_id or "aws/opensearch" in kms_key_id:
             return f"{kb_name} (OpenSearch domain not using CMK)"
     except ClientError as e:
         LOGGER.error(f"Error checking OpenSearch domain: {str(e)}")
         return f"{kb_name} (error checking OpenSearch domain)"
     return None
 
 
-def check_knowledge_base(kb_id: str, kb_name: str) -> str | None:  # type: ignore  # noqa: CFQ004
+def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str | None]:  # type: ignore  # noqa: CFQ004
     """Check a knowledge base's OpenSearch configuration.
 
     Args:
@@ -91,40 +92,56 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> str | None:  # type: ignor
         ClientError: If there is an error checking the knowledge base
 
     Returns:
-        str | None: Error message if non-compliant, None if compliant
+        tuple[bool, str | None]: (has_opensearch, error_message)
     """
     try:
         kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id)
-        vector_store = kb_details.get("vectorStoreConfiguration")
+        # Convert datetime objects to strings before JSON serialization
+        kb_details_serializable = json.loads(json.dumps(kb_details, default=str))
+        LOGGER.info(f"Knowledge base details for {kb_name}: {json.dumps(kb_details_serializable)}")
+
+        # Access the knowledgeBase key from the response
+        kb_data = kb_details.get("knowledgeBase", {})
+
+        # Check both possible locations for vector store config
+        vector_store = kb_data.get("vectorStoreConfiguration") or kb_data.get("storageConfiguration", {})
+        LOGGER.info(f"Vector store config for {kb_name}: {json.dumps(vector_store)}")
 
         if not vector_store or not isinstance(vector_store, dict):
-            return None
+            LOGGER.info(f"No vector store configuration found for {kb_name}")
+            return False, None
 
-        if vector_store.get("vectorStoreType") != "OPENSEARCH":
-            return None
+        vector_store_type = vector_store.get("vectorStoreType") or vector_store.get("type")
+        LOGGER.info(f"Vector store type for {kb_name}: {vector_store_type}")
+        if not vector_store_type or (vector_store_type.upper() != "OPENSEARCH" and vector_store_type.upper() != "OPENSEARCH_SERVERLESS"):
+            LOGGER.info(f"Vector store type is not OpenSearch for {kb_name}")
+            return False, None
 
         opensearch_config = vector_store.get("opensearchServerlessConfiguration") or vector_store.get("opensearchConfiguration")
+        LOGGER.info(f"OpenSearch config for {kb_name}: {json.dumps(opensearch_config)}")
         if not opensearch_config:
-            return f"{kb_name} (missing OpenSearch configuration)"
+            return True, f"{kb_name} (missing OpenSearch configuration)"
 
         if "collectionArn" in opensearch_config:
             collection_id = opensearch_config["collectionArn"].split("/")[-1]
-            return check_opensearch_serverless(collection_id, kb_name)
+            LOGGER.info(f"Found OpenSearch Serverless collection {collection_id} for {kb_name}")
+            return True, check_opensearch_serverless(collection_id, kb_name)
 
         domain_endpoint = opensearch_config.get("endpoint", "")
         if not domain_endpoint:
-            return f"{kb_name} (missing OpenSearch domain endpoint)"
+            return True, f"{kb_name} (missing OpenSearch domain endpoint)"
         domain_name = domain_endpoint.split(".")[0]
-        return check_opensearch_domain(domain_name, kb_name)
+        LOGGER.info(f"Found OpenSearch domain {domain_name} for {kb_name}")
+        return True, check_opensearch_domain(domain_name, kb_name)
 
     except ClientError as e:
         LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}")
         if e.response["Error"]["Code"] == "AccessDeniedException":
-            return f"{kb_name} (access denied)"
+            return True, f"{kb_name} (access denied)"
         raise
 
 
-def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
+def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100, CFQ004
     """Evaluate if Bedrock Knowledge Base OpenSearch vector stores are encrypted with KMS CMK.
 
     Args:
@@ -135,16 +152,20 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
     """
     try:
         non_compliant_kbs = []
+        has_opensearch = False
         paginator = bedrock_agent_client.get_paginator("list_knowledge_bases")
 
         for page in paginator.paginate():
             for kb in page["knowledgeBaseSummaries"]:
                 kb_id = kb["knowledgeBaseId"]
                 kb_name = kb.get("name", kb_id)
-                error = check_knowledge_base(kb_id, kb_name)
+                is_opensearch, error = check_knowledge_base(kb_id, kb_name)
+                has_opensearch = has_opensearch or is_opensearch
                 if error:
                     non_compliant_kbs.append(error)
 
+        if not has_opensearch:
+            return "COMPLIANT", "No OpenSearch vector stores found in knowledge bases"
         if non_compliant_kbs:
             return "NON_COMPLIANT", (
                 "The following knowledge bases have OpenSearch vector stores not encrypted with CMK: " + f"{'; '.join(non_compliant_kbs)}"
@@ -153,7 +174,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
 
     except Exception as e:
         LOGGER.error(f"Error evaluating Bedrock Knowledge Base OpenSearch encryption: {str(e)}")
-        return "ERROR", f"Error evaluating compliance: {str(e)}"
+        return "INSUFFICIENT_DATA", f"Error evaluating compliance: {str(e)}"
 
 
 def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100

aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json

@@ -224,6 +224,12 @@
           "opensearch:GetSecurityPolicy"
         ],
         "Resource": "*"
+      },
+      {
+        "Sid": "AllowOpenSearchServerlessAccess",
+        "Effect": "Allow",
+        "Action": ["aoss:GetSecurityPolicy"],
+        "Resource": "*"
       }
     ]
   }