SRA GenAI Deep-dive capability two controls (#292)

* adding prototype kb logging check rule

* fix config rule

* adding kb ingestion encryption check rule

* added kb s3 bucket check rule

* added prototype config rule for kb vector store secrete

* added check for opensearch encryption (if used in KB)

* update param name

* working to fix bug

* updating lambda code. updating permissions

* return arn if already exists

* update README.md file

* update readme

* updating for flake8 and mypy errors

* fixing more flake8 and mypy errors

* fixing formatter errors

* update pyproject.toml

* update pyproject.toml

* fixing errors and permissions

* refactor for mypy fixes; updates for finding kmsarn in collection for kb

* fixing for changes made because of mypy

* updating for mypy fixes; also resolved API call issues

* flake8 issue

* fixes from flake8 and mypy refactoring; update api calls for failures

* fixes for refactoring due to mypy and flake8; update API calls to appropriately get info

* fix annotation language

* black formatter changes

* update perms

* ensuring annotations do not exceed 256 char limit

* black formatter changes

* ensuring rule delete operations don't target accounts not in rule_accounts (e.g. mgmt acct)

* adding error handling with IAM policy deletes (race condition found)

* handling flake8 issues

* fixing cleanup issues

* don't record CW centrol observability resources in dry-run if deploy is false

* update condition statement

* enable key rotation for CMKs

---------

Co-authored-by: liamschn <57731583+liamschn@users.noreply.github.com>

Author: cyphronix

aws_sra_examples/solutions/genai/bedrock_org/README.md

@@ -7,6 +7,7 @@
 - [Security Controls](#security-controls)
 - [JSON Parameters](#json-parameters)
 - [References](#references)
+- [Related Security Control Solutions](#related-security-control-solutions)
 
 ---
 
@@ -102,6 +103,11 @@ aws cloudformation create-stack \
         ParameterKey=pBedrockPromptInjectionFilterParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\"], \"filter_params\": {\"log_group_name\": \"model-invocation-log-group\", \"input_path\": \"input.inputBodyJson.messages[0].content\"}}"' \
         ParameterKey=pBedrockSensitiveInfoFilterParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\"], \"filter_params\": {\"log_group_name\": \"model-invocation-log-group\", \"input_path\": \"input.inputBodyJson.messages[0].content\"}}"' \
         ParameterKey=pBedrockCentralObservabilityParams,ParameterValue='"{\"deploy\": \"true\", \"bedrock_accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\"]}"' \
+        ParameterKey=pBedrockKBLoggingRuleParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\",\"us-west-2\"], \"input_params\": {}}"' \
+        ParameterKey=pBedrockKBIngestionEncryptionRuleParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\",\"us-west-2\"], \"input_params\": {}}"' \
+        ParameterKey=pBedrockKBS3BucketRuleParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\",\"us-west-2\"], \"input_params\": {\"check_retention\": \"true\", \"check_encryption\": \"true\", \"check_access_logging\": \"true\", \"check_object_locking\": \"true\", \"check_versioning\": \"true\"}}"' \
+        ParameterKey=pBedrockKBVectorStoreSecretRuleParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\",\"us-west-2\"], \"input_params\": {}}"' \
+        ParameterKey=pBedrockKBOpenSearchEncryptionRuleParams,ParameterValue='"{\"deploy\": \"true\", \"accounts\": [\"222222222222\",\"333333333333\"], \"regions\": [\"us-east-1\",\"us-west-2\"], \"input_params\": {}}"' \
     --capabilities CAPABILITY_NAMED_IAM
 ```
 
@@ -139,6 +145,11 @@ Please read the following notes before deploying the stack to ensure successful
 | CloudWatch Endpoint Validation | Ensures proper CloudWatch VPC endpoint setup | [pBedrockCWEndpointsRuleParams](#pbedrockcwendpointsruleparams) |
 | S3 Endpoint Validation | Ensures proper S3 VPC endpoint setup | [pBedrockS3EndpointsRuleParams](#pbedrocks3endpointsruleparams) |
 | Guardrail Encryption | Validates KMS encryption for Bedrock guardrails | [pBedrockGuardrailEncryptionRuleParams](#pbedrockguardrailencryptionruleparams) |
+| Knowledge Base Logging | Validates logging configuration for Bedrock Knowledge Base | [pBedrockKBLoggingRuleParams](#pbedrockkbloggingruleparams) |
+| Knowledge Base Ingestion Encryption | Validates encryption for Knowledge Base data ingestion | [pBedrockKBIngestionEncryptionRuleParams](#pbedrockkbingestionencryptionruleparams) |
+| Knowledge Base S3 Bucket | Validates S3 bucket configurations for Knowledge Base | [pBedrockKBS3BucketRuleParams](#pbedrockkbs3bucketruleparams) |
+| Knowledge Base Vector Store Secret | Validates vector store secret configuration | [pBedrockKBVectorStoreSecretRuleParams](#pbedrockkbvectorstoresecretruleparams) |
+| Knowledge Base OpenSearch Encryption | Validates OpenSearch encryption configuration | [pBedrockKBOpenSearchEncryptionRuleParams](#pbedrockkbopensearchencryptionruleparams) |
 
 > **Important Note**: The Config rule Lambda execution role needs to have access to any KMS keys used to encrypt Bedrock guardrails. Make sure to grant the appropriate KMS key permissions to the Lambda role to ensure proper evaluation of encrypted guardrail configurations.
 
@@ -155,6 +166,15 @@ Please read the following notes before deploying the stack to ensure successful
 |-----------------|-------------|----------------|
 | Central Observability | Configures cross-account/region metric aggregation | [pBedrockCentralObservabilityParams](#pbedrockcentralobservabilityparams) |
 
+### Bedrock Knowledge Base
+| Security Control | Description | JSON Parameter |
+|-----------------|-------------|----------------|
+| KB Logging | Validates logging configuration for Bedrock Knowledge Base | [pBedrockKBLoggingRuleParams](#pbedrockkbloggingruleparams) |
+| KB Ingestion Encryption | Validates encryption configuration for Bedrock Knowledge Base | [pBedrockKBIngestionEncryptionRuleParams](#pbedrockkbingestionencryptionruleparams) |
+| KB S3 Bucket | Validates S3 bucket configuration for Bedrock Knowledge Base | [pBedrockKBS3BucketRuleParams](#pbedrockkbs3bucketruleparams) |
+| KB Vector Store Secret | Validates secret configuration for Bedrock Knowledge Base | [pBedrockKBVectorStoreSecretRuleParams](#pbedrockkbvectorstoresecretruleparams) |
+| KB OpenSearch Encryption | Validates encryption configuration for Bedrock Knowledge Base | [pBedrockKBOpenSearchEncryptionRuleParams](#pbedrockkbopensearchencryptionruleparams) |
+
 ---
 ## JSON Parameters
 
@@ -367,6 +387,72 @@ This section explains the parameters in the CloudFormation template that require
 }
 ```
 
+### `pBedrockKBLoggingRuleParams`
+- **Purpose**: Validates logging configuration for Bedrock Knowledge Base.
+- **Structure**:
+```json
+{
+  "deploy": "true|false",
+  "accounts": ["account_id1", "account_id2"],
+  "regions": ["region1", "region2"],
+  "input_params": {}
+}
+```
+
+### `pBedrockKBIngestionEncryptionRuleParams`
+- **Purpose**: Validates encryption configuration for Bedrock Knowledge Base.
+- **Structure**:
+```json
+{
+  "deploy": "true|false",
+  "accounts": ["account_id1", "account_id2"],
+  "regions": ["region1", "region2"],
+  "input_params": {}
+}
+```
+
+### `pBedrockKBS3BucketRuleParams`
+- **Purpose**: Validates S3 bucket configuration for Bedrock Knowledge Base.
+- **Structure**:
+```json
+{
+  "deploy": "true|false",
+  "accounts": ["account_id1", "account_id2"],
+  "regions": ["region1", "region2"],
+  "input_params": {
+    "check_retention": "true|false",
+    "check_encryption": "true|false",
+    "check_access_logging": "true|false",
+    "check_object_locking": "true|false",
+    "check_versioning": "true|false"
+  }
+}
+```
+
+### `pBedrockKBVectorStoreSecretRuleParams`
+- **Purpose**: Validates secret configuration for Bedrock Knowledge Base.
+- **Structure**:
+```json
+{
+  "deploy": "true|false",
+  "accounts": ["account_id1", "account_id2"],
+  "regions": ["region1", "region2"],
+  "input_params": {}
+}
+```
+
+### `pBedrockKBOpenSearchEncryptionRuleParams`
+- **Purpose**: Validates encryption configuration for Bedrock Knowledge Base.
+- **Structure**:
+```json
+{
+  "deploy": "true|false",
+  "accounts": ["account_id1", "account_id2"],
+  "regions": ["region1", "region2"],
+  "input_params": {}
+}
+```
+
 ---
 ## References
 - [AWS SRA Generative AI Deep-Dive](https://docs.aws.amazon.com/prescriptive-guidance/latest/security-reference-architecture/gen-ai-sra.html)
@@ -375,3 +461,32 @@ This section explains the parameters in the CloudFormation template that require
 - [CloudWatch Metrics and Alarms](https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/WhatIsCloudWatch.html)
 - [AWS Lambda](https://docs.aws.amazon.com/lambda/latest/dg/welcome.html)
 - [AWS KMS](https://docs.aws.amazon.com/kms/latest/developerguide/overview.html)
+
+## Related Security Control Solutions
+
+This solution works in conjunction with other AWS SRA solutions to provide comprehensive security controls for Bedrock GenAI environments:
+
+### Amazon Bedrock Guardrails Solution
+The [SRA Bedrock Guardrails solution](../../genai/bedrock_guardrails/README.md) provides automated deployment of Amazon Bedrock Guardrails across your organization. It supports:
+
+- **Content Filters**: Block harmful content in inputs/outputs based on predefined categories (Hate, Insults, Sexual, Violence, Misconduct, Prompt Attack)
+- **Denied Topics**: Define and block undesirable topics
+- **Word Filters**: Block specific words, phrases, and profanity
+- **Sensitive Information Filters**: Block or mask PII and sensitive data
+- **Contextual Grounding**: Detect and filter hallucinations based on source grounding
+
+The solution uses KMS encryption for enhanced security and requires proper IAM role configurations for users who need to invoke or manage guardrails.
+
+### GuardDuty Malware Protection for S3
+The [SRA GuardDuty Malware Protection solution](../../guardduty/guardduty_malware_protection_for_s3/README.md) helps protect S3 buckets used in your Bedrock environment from malware. This is particularly important for:
+
+- Model evaluation job buckets
+- Knowledge base data ingestion buckets
+- Model invocation logging buckets
+
+The solution enables GuardDuty's malware scanning capabilities to detect malicious files that could be used in prompt injection attacks or compromise your GenAI applications.
+
+These complementary solutions work together to provide defense-in-depth for your Bedrock GenAI environment:
+- This solution (SRA Bedrock Org) provides organizational security controls and monitoring
+- Bedrock Guardrails solution provides content and data security controls
+- GuardDuty Malware Protection ensures S3 bucket security against malware threats

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import json
 import logging
 import os

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import ast
 import logging
 import os

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import json
 import logging
 import os

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import ast
 import json
 import logging

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import json
 import logging
 import os

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import json
 import logging
 import os

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

@@ -7,6 +7,7 @@
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
 """
+
 import json
 import logging
 import os

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

@@ -0,0 +1,151 @@
+"""Config rule to check knowledge base data ingestion encryption for Bedrock environments.
+
+Version: 1.0
+
+Config rule for SRA in the repo, https://github.com/aws-samples/aws-security-reference-architecture-examples
+
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: MIT-0
+"""
+
+import json
+import logging
+import os
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+# Setup Default Logger
+LOGGER = logging.getLogger(__name__)
+log_level = os.environ.get("LOG_LEVEL", logging.INFO)
+LOGGER.setLevel(log_level)
+LOGGER.info(f"boto3 version: {boto3.__version__}")
+
+# Get AWS region from environment variable
+AWS_REGION = os.environ.get("AWS_REGION")
+
+# Initialize AWS clients
+bedrock_agent_client = boto3.client("bedrock-agent", region_name=AWS_REGION)
+config_client = boto3.client("config", region_name=AWS_REGION)
+
+
+def check_data_sources(kb_id: str, kb_name: str) -> str | None:  # type: ignore  # noqa: CFQ004, CCR001
+    """Check if a knowledge base's data sources are encrypted with KMS during ingestion.
+
+    Args:
+        kb_id (str): Knowledge base ID
+        kb_name (str): Knowledge base name
+
+    Raises:
+        ClientError: If there is an error checking the knowledge base
+
+    Returns:
+        str | None: Error message if non-compliant, None if compliant
+    """
+    try:
+        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 response"
+
+        unencrypted_sources = []
+        for source in data_sources.get("dataSourceSummaries", []):
+            LOGGER.info(f"Source: {source}")
+            if not isinstance(source, dict):
+                continue
+
+            # Get the detailed data source configuration
+            try:
+                source_details = bedrock_agent_client.get_data_source(knowledgeBaseId=kb_id, dataSourceId=source["dataSourceId"])
+                LOGGER.info(f"Source details: {source_details}")
+
+                # Check for KMS encryption configuration
+                data_source = source_details.get("dataSource", {})
+                encryption_config = data_source.get("serverSideEncryptionConfiguration", {})
+                LOGGER.info(f"Encryption config: {encryption_config}")
+
+                # Check if KMS key is configured for encryption
+                if not encryption_config.get("kmsKeyArn"):
+                    unencrypted_sources.append(source.get("name", source["dataSourceId"]))
+
+            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'])}")
+                continue
+
+        if 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"
+        raise
+
+
+def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]:  # noqa: U100
+    """Evaluate if Bedrock Knowledge Base data sources are encrypted with KMS.
+
+    Args:
+        rule_parameters (dict): Rule parameters from AWS Config rule.
+
+    Returns:
+        tuple[str, str]: Compliance type and annotation message.
+    """
+    try:
+        non_compliant_kbs = []
+        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_data_sources(kb_id, kb_name)
+                if error:
+                    non_compliant_kbs.append(error)
+
+        if 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 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: {str(e)[:240]}"
+
+
+def lambda_handler(event: dict, context: Any) -> None:  # noqa: U100
+    """Lambda handler.
+
+    Args:
+        event (dict): Lambda event object
+        context (Any): Lambda context object
+    """
+    LOGGER.info("Evaluating compliance for AWS Config rule")
+    LOGGER.info(f"Event: {json.dumps(event)}")
+
+    invoking_event = json.loads(event["invokingEvent"])
+    rule_parameters = json.loads(event["ruleParameters"]) if "ruleParameters" in event else {}
+
+    compliance_type, annotation = evaluate_compliance(rule_parameters)
+
+    evaluation = {
+        "ComplianceResourceType": "AWS::::Account",
+        "ComplianceResourceId": event["accountId"],
+        "ComplianceType": compliance_type,
+        "Annotation": annotation,
+        "OrderingTimestamp": invoking_event["notificationCreationTime"],
+    }
+
+    LOGGER.info(f"Compliance evaluation result: {compliance_type}")
+    LOGGER.info(f"Annotation: {annotation}")
+
+    config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"])  # type: ignore
+
+    LOGGER.info("Compliance evaluation complete.")