ensuring annotations do not exceed 256 char limit

Author: cyphronix

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

@@ -47,7 +47,7 @@ def check_data_sources(kb_id: str, kb_name: str) -> str | None:  # type: ignore
         data_sources = bedrock_agent_client.list_data_sources(knowledgeBaseId=kb_id)
         LOGGER.info(f"Data sources: {data_sources}")
         if not isinstance(data_sources, dict):
-            return f"{kb_name} (invalid data sources response)"
+            return f"{kb_name}: Invalid response"
 
         unencrypted_sources = []
         for source in data_sources.get("dataSourceSummaries", []):
@@ -72,16 +72,16 @@ def check_data_sources(kb_id: str, kb_name: str) -> str | None:  # type: ignore
             except ClientError as e:
                 LOGGER.error(f"Error getting data source details for {source.get('name', source['dataSourceId'])}: {str(e)}")
                 if e.response["Error"]["Code"] == "AccessDeniedException":
-                    unencrypted_sources.append(f"{source.get('name', source['dataSourceId'])} (access denied)")
+                    unencrypted_sources.append(f"{source.get('name', source['dataSourceId'])}")
                 continue
 
         if unencrypted_sources:
-            return f"{kb_name} (sources using default AWS-managed key instead of Customer Managed Key: {', '.join(unencrypted_sources)})"
+            return f"{kb_name}: {len(unencrypted_sources)} sources need CMK"
         return None
     except ClientError as e:
         LOGGER.error(f"Error checking data sources for knowledge base {kb_name}: {str(e)}")
         if e.response["Error"]["Code"] == "AccessDeniedException":
-            return f"{kb_name} (access denied)"
+            return f"{kb_name}: Access denied"
         raise
 
 
@@ -107,16 +107,17 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
                     non_compliant_kbs.append(error)
 
         if non_compliant_kbs:
-            msg = (
-                "The following knowledge bases are using default AWS-managed keys "
-                + f"instead of Customer Managed Keys: {'; '.join(non_compliant_kbs)}"
-            )
+            msg = f"KBs missing Customer Managed Keys: {'; '.join(non_compliant_kbs)}"
+            # Ensure annotation doesn't exceed 256 characters
+            if len(msg) > 256:
+                LOGGER.info(f"Full message truncated: {msg}")
+                msg = msg[:220] + " (see CloudWatch logs for details)"
             return "NON_COMPLIANT", msg
-        return "COMPLIANT", "All knowledge base data sources are encrypted with Customer Managed Keys"
+        return "COMPLIANT", "All KB data sources use Customer Managed Keys"
 
     except Exception as e:
         LOGGER.error(f"Error evaluating Bedrock Knowledge Base encryption: {str(e)}")
-        return "ERROR", f"Error evaluating compliance: {str(e)}"
+        return "ERROR", f"Error: {str(e)[:240]}"
 
 
 def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100

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

@@ -30,6 +30,32 @@
 logs_client = boto3.client("logs", region_name=AWS_REGION)
 sts_client = boto3.client("sts", region_name=AWS_REGION)
 
+# Max length for AWS Config annotation
+MAX_ANNOTATION_LENGTH = 256
+
+
+def truncate_annotation(message: str) -> str:
+    """Ensure annotation stays within AWS Config's 256 character limit.
+
+    Args:
+        message (str): Original annotation message
+
+    Returns:
+        str: Truncated message with CloudWatch reference if needed
+    """
+    if len(message) <= MAX_ANNOTATION_LENGTH:
+        return message
+
+    log_group = f"/aws/lambda/{os.environ.get('AWS_LAMBDA_FUNCTION_NAME', 'unknown')}"
+    reference = f" See CloudWatch logs ({log_group}) for details."
+
+    # Calculate available space for the actual message
+    available_chars = MAX_ANNOTATION_LENGTH - len(reference)
+
+    # Truncate message and add reference
+    truncated = message[:available_chars - 3] + "..."
+    return truncated + reference
+
 
 def check_kb_logging(kb_id: str) -> Tuple[bool, Optional[str]]:  # noqa: CCR001
     """Check if knowledge base has CloudWatch logging enabled.
@@ -117,10 +143,10 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: CFQ0
             kb_list.extend(page.get("knowledgeBaseSummaries", []))
 
         if not kb_list:
-            return "COMPLIANT", "No knowledge bases found in the account"
+            return "COMPLIANT", "No KBs found"
 
         non_compliant_kbs = []
-        compliant_kbs = []
+        compliant_count = 0
 
         # Check each knowledge base for logging configuration
         for kb in kb_list:
@@ -129,17 +155,24 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: CFQ0
 
             has_logging, destination_type = check_kb_logging(kb_id)
             if not has_logging:
-                non_compliant_kbs.append(f"{kb_id} ({kb_name}) - logging not configured")
+                # Use shorter format for non-compliant KBs
+                non_compliant_kbs.append(f"{kb_id[:8]}..({kb_name[:10]})")
             else:
-                compliant_kbs.append(f"{kb_id} ({kb_name}) - logging configured to {destination_type}")
+                compliant_count += 1
+                LOGGER.info(f"KB {kb_id} ({kb_name}) has logging to {destination_type}")
 
         if non_compliant_kbs:
-            return "NON_COMPLIANT", f"The following knowledge bases do not have logging enabled: {', '.join(non_compliant_kbs)}"
-        return "COMPLIANT", f"All knowledge bases have logging enabled: {', '.join(compliant_kbs)}"
+            msg = f"{len(non_compliant_kbs)} KBs without logging: {', '.join(non_compliant_kbs[:5])}"
+            # Add count indicator if there are more than shown
+            if len(non_compliant_kbs) > 5:
+                msg += f" +{len(non_compliant_kbs) - 5} more"
+            return "NON_COMPLIANT", truncate_annotation(msg)
+
+        return "COMPLIANT", truncate_annotation(f"All {compliant_count} KBs have logging enabled")
 
     except Exception as e:
-        LOGGER.error(f"Error evaluating Bedrock Knowledge Base logging configuration: {str(e)}")
-        return "ERROR", f"Error evaluating compliance: {str(e)}"
+        LOGGER.error(f"Error evaluating Bedrock KB logging: {str(e)}")
+        return "ERROR", truncate_annotation(f"Error: {str(e)}")
 
 
 def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100

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

@@ -49,12 +49,12 @@ def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None:
 
         if not collection_response.get("collectionDetails"):
             LOGGER.error(f"No collection details found for ID {collection_id}")
-            return f"{kb_name} (OpenSearch Serverless collection not found)"
+            return f"{kb_name} (no collection)"
 
         collection_name = collection_response["collectionDetails"][0].get("name")
         if not collection_name:
             LOGGER.error(f"No collection name found for ID {collection_id}")
-            return f"{kb_name} (OpenSearch Serverless collection name not found)"
+            return f"{kb_name} (no collection name)"
 
         # Get the specific policy details using the collection name
         policy_details = opensearch_serverless_client.get_security_policy(name=collection_name, type="encryption")
@@ -65,19 +65,19 @@ def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None:
         LOGGER.info(f"Policy details dict (after getting policy): {json.dumps(policy_details_dict, default=str)}")
 
         if policy_details_dict.get("AWSOwnedKey", False):
-            LOGGER.info(f"{kb_name} (OpenSearch Serverless using AWS-owned key instead of CMK)")
-            return f"{kb_name} (OpenSearch Serverless using AWS-owned key instead of CMK)"
+            LOGGER.info(f"{kb_name} (Using AWS-owned key, not CMK)")
+            return f"{kb_name} (AWS-owned key)"
 
         kms_key_arn = policy_details_dict.get("KmsARN", "")
         if not kms_key_arn:
             LOGGER.info(f"{kb_name} (OpenSearch Serverless not using CMK)")
-            return f"{kb_name} (OpenSearch Serverless not using CMK)"
+            return f"{kb_name} (no CMK)"
 
         return None
 
     except ClientError as e:
         LOGGER.error(f"Error checking OpenSearch Serverless collection: {str(e)}")
-        return f"{kb_name} (error checking OpenSearch Serverless)"
+        return f"{kb_name} (error)"
 
 
 def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None:  # type: ignore  # noqa: CFQ004
@@ -94,13 +94,13 @@ def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None:  # ty
         domain = opensearch_client.describe_domain(DomainName=domain_name)
         encryption_config = domain.get("DomainStatus", {}).get("EncryptionAtRestOptions", {})
         if not encryption_config.get("Enabled", False):
-            return f"{kb_name} (OpenSearch domain encryption not enabled)"
+            return f"{kb_name} (encryption disabled)"
         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)"
+            return f"{kb_name} (no CMK)"
     except ClientError as e:
         LOGGER.error(f"Error checking OpenSearch domain: {str(e)}")
-        return f"{kb_name} (error checking OpenSearch domain)"
+        return f"{kb_name} (error)"
     return None
 
 
@@ -143,7 +143,7 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str | 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 True, f"{kb_name} (missing OpenSearch configuration)"
+            return True, f"{kb_name} (missing config)"
 
         if "collectionArn" in opensearch_config:
             collection_id = opensearch_config["collectionArn"].split("/")[-1]
@@ -152,7 +152,7 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str | None]:
 
         domain_endpoint = opensearch_config.get("endpoint", "")
         if not domain_endpoint:
-            return True, f"{kb_name} (missing OpenSearch domain endpoint)"
+            return True, f"{kb_name} (no endpoint)"
         domain_name = domain_endpoint.split(".")[0]
         LOGGER.info(f"Found OpenSearch domain {domain_name} for {kb_name}")
         return True, check_opensearch_domain(domain_name, kb_name)
@@ -164,11 +164,12 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str | None]:
         raise
 
 
-def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100, CFQ004
+def evaluate_compliance(rule_parameters: dict, request_id: str = "") -> tuple[str, str]:  # noqa: U100, CFQ004
     """Evaluate if Bedrock Knowledge Base OpenSearch vector stores are encrypted with KMS CMK.
 
     Args:
         rule_parameters (dict): Rule parameters from AWS Config rule.
+        request_id (str): Lambda request ID for CloudWatch log reference.
 
     Returns:
         tuple[str, str]: Compliance type and annotation message.
@@ -188,16 +189,21 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
                     non_compliant_kbs.append(error)
 
         if not has_opensearch:
-            return "COMPLIANT", "No OpenSearch vector stores found in knowledge bases"
+            return "COMPLIANT", "No OpenSearch vector stores found"
+
         if non_compliant_kbs:
-            return "NON_COMPLIANT", (
-                "The following knowledge bases have OpenSearch vector stores not encrypted with CMK: " + f"{'; '.join(non_compliant_kbs)}"
-            )
-        return "COMPLIANT", "All knowledge base OpenSearch vector stores are encrypted with KMS CMK"
+            message = "KBs without CMK encryption: " + ", ".join(non_compliant_kbs)
+            # Check if message exceeds the 256-character limit
+            if len(message) > 256:
+                LOGGER.info(f"Full message (truncated in annotation): {message}")
+                return "NON_COMPLIANT", f"Multiple KBs without CMK encryption. See CloudWatch logs ({request_id})"
+            return "NON_COMPLIANT", message
+
+        return "COMPLIANT", "All KBs properly encrypted with CMK"
 
     except Exception as e:
         LOGGER.error(f"Error evaluating Bedrock Knowledge Base OpenSearch encryption: {str(e)}")
-        return "INSUFFICIENT_DATA", f"Error evaluating compliance: {str(e)}"
+        return "INSUFFICIENT_DATA", f"Error: {str(e)[:220]}"
 
 
 def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100
@@ -209,11 +215,17 @@ def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100
     """
     LOGGER.info("Evaluating compliance for AWS Config rule")
     LOGGER.info(f"Event: {json.dumps(event)}")
+    LOGGER.info(f"Lambda Request ID: {context.aws_request_id}")
 
     invoking_event = json.loads(event["invokingEvent"])
     rule_parameters = json.loads(event["ruleParameters"]) if "ruleParameters" in event else {}
 
-    compliance_type, annotation = evaluate_compliance(rule_parameters)
+    compliance_type, annotation = evaluate_compliance(rule_parameters, context.aws_request_id)
+
+    # Ensure annotation doesn't exceed 256 characters
+    if len(annotation) > 256:
+        LOGGER.info(f"Original annotation (truncated): {annotation}")
+        annotation = annotation[:252] + "..."
 
     evaluation = {
         "ComplianceResourceType": "AWS::::Account",

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

@@ -223,12 +223,35 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:
                 non_compliant_buckets.extend(check_knowledge_base(kb["knowledgeBaseId"], rule_parameters))
 
         if non_compliant_buckets:
-            return "NON_COMPLIANT", f"The following KB S3 buckets are non-compliant: {'; '.join(non_compliant_buckets)}"
-        return "COMPLIANT", "All Knowledge Base S3 buckets meet the required configurations"
+            # Create a shorter message for each bucket by using abbreviations
+            bucket_msgs = []
+            for bucket in non_compliant_buckets:
+                # Replace longer descriptions with abbreviations
+                short_msg = bucket.replace("missing: ", "")
+                short_msg = short_msg.replace("retention", "ret")
+                short_msg = short_msg.replace("encryption", "enc")
+                short_msg = short_msg.replace("access logging", "log")
+                short_msg = short_msg.replace("object locking", "lock")
+                short_msg = short_msg.replace("versioning", "ver")
+                bucket_msgs.append(short_msg)
+
+            # Build the annotation message
+            annotation = f"Non-compliant KB S3 buckets: {'; '.join(bucket_msgs)}"
+
+            # If annotation exceeds limit, truncate and refer to logs
+            if len(annotation) > 256:
+                # Log the full message
+                LOGGER.info(f"Full compliance details: {annotation}")
+                # Create a truncated message that fits within the limit
+                count = len(non_compliant_buckets)
+                annotation = f"{count} non-compliant KB S3 buckets. See CloudWatch logs for details."
+
+            return "NON_COMPLIANT", annotation
+        return "COMPLIANT", "All KB S3 buckets compliant"
 
     except Exception as e:
         LOGGER.error(f"Error evaluating Knowledge Base S3 bucket configurations: {str(e)}")
-        return "ERROR", f"Error evaluating compliance: {str(e)}"
+        return "ERROR", f"Error: {str(e)[:240]}"
 
 
 def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100