From 8fb8ae7e32bd001abd72d0d14c1c5b41804c5625 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Feb 2025 09:59:38 -0700 Subject: [PATCH 01/36] adding prototype kb logging check rule --- .../rules/sra_bedrock_check_kb_logging/app.py | 110 ++++++++++++++++++ .../genai/bedrock_org/lambda/src/app.py | 2 + .../sra_config_lambda_iam_permissions.json | 14 +++ .../templates/sra-bedrock-org-main.yaml | 16 ++- 4 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py new file mode 100644 index 000000000..53e67c11e --- /dev/null +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -0,0 +1,110 @@ +"""Config rule to check knowledge base logging 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_client = boto3.client("bedrock", region_name=AWS_REGION) +config_client = boto3.client("config", region_name=AWS_REGION) + + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + """Evaluate if Bedrock Knowledge Base logging is properly configured. + + Args: + rule_parameters (dict): Rule parameters from AWS Config rule. + + Returns: + tuple[str, str]: Compliance type and annotation message. + """ + try: + # List all knowledge bases + kb_list = [] + paginator = bedrock_client.get_paginator('list_knowledge_bases') + for page in paginator.paginate(): + kb_list.extend(page.get('knowledgeBases', [])) + + if not kb_list: + return "COMPLIANT", "No knowledge bases found in the account" + + non_compliant_kbs = [] + + # Check each knowledge base for logging configuration + for kb in kb_list: + kb_id = kb['knowledgeBaseId'] + try: + kb_details = bedrock_client.get_knowledge_base( + knowledgeBaseId=kb_id + ) + + # Check if logging is enabled + logging_config = kb_details.get('loggingConfiguration', {}) + if not logging_config or not logging_config.get('enabled', False): + non_compliant_kbs.append(f"{kb_id} ({kb.get('name', 'unnamed')})") + + except ClientError as e: + LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") + if e.response['Error']['Code'] == 'AccessDeniedException': + non_compliant_kbs.append(f"{kb_id} (access denied)") + else: + raise + + if non_compliant_kbs: + return "NON_COMPLIANT", f"The following knowledge bases do not have logging enabled: {', '.join(non_compliant_kbs)}" + return "COMPLIANT", "All knowledge bases 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)}" + + +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.") \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 17d2edc1c..db2ff6c25 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -200,6 +200,8 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"filter_params"\s*:\s*\{"log_group_name"\s*:\s*"[^"\s]+",\s*"input_path"\s*:\s*"[^"\s]+"\}\}$', "SRA-BEDROCK-CENTRAL-OBSERVABILITY": r'^\{"deploy"\s*:\s*"(true|false)",\s*"bedrock_accounts"\s*:\s*' + r'\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\]\}$', + "SRA-BEDROCK-CHECK-KB-LOGGING": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', } # Instantiate sra class objects diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index de11d6ac8..8d829ae61 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -123,5 +123,19 @@ "Resource": "*" } ] + }, + "sra-bedrock-check-kb-logging": { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "AllowKnowledgeBaseAccess", + "Effect": "Allow", + "Action": [ + "bedrock:ListKnowledgeBases", + "bedrock:GetKnowledgeBase" + ], + "Resource": "*" + } + ] } } diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index 1af2e701c..2ea8fc31a 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -256,7 +256,6 @@ Parameters: or for titan: {"deploy": "true", "filter_params": {"log_group_name": "model-invocation-log-group", "input_path": "input.inputBodyJson.inputText"}} NOTE: input_path is based on the base model used such as clause or titan; check the invocation log InvokeModel messages for details - pBedrockCentralObservabilityParams: Type: String Default: '{"deploy": "true", "bedrock_accounts": ["444455556666"], "regions": ["us-west-2"]}' @@ -282,6 +281,17 @@ Parameters: ConstraintDescription: > Must be a valid JSON string containing an array of region names. Example: ["us-east-1", "us-west-2"] + pBedrockKBLoggingRuleParams: + Type: String + Default: '{"deploy": "true", "accounts": ["444455556666"], "regions": ["us-west-2"], "input_params": {}}' + Description: Bedrock Knowledge Base Logging Config Rule Parameters + AllowedPattern: ^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$ + ConstraintDescription: + "Must be a valid JSON string containing: 'deploy' (true/false), 'accounts' (array of account numbers), + 'regions' (array of region names), and 'input_params' object/dict (input params must be empty). Arrays can be empty. + Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or + {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + Metadata: AWS::CloudFormation::Interface: ParameterGroups: @@ -322,6 +332,7 @@ Metadata: - pBedrockCWEndpointsRuleParams - pBedrockS3EndpointsRuleParams - pBedrockGuardrailEncryptionRuleParams + - pBedrockKBLoggingRuleParams - Label: default: Bedrock CloudWatch Metric Filters Parameters: @@ -389,6 +400,8 @@ Metadata: default: Bedrock Accounts pBedrockRegions: default: Bedrock Regions + pBedrockKBLoggingRuleParams: + default: Bedrock Knowledge Base Logging Config Rule Parameters Resources: rBedrockOrgLambdaRole: @@ -669,6 +682,7 @@ Resources: SRA-BEDROCK-FILTER-PROMPT-INJECTION: !Ref pBedrockPromptInjectionFilterParams SRA-BEDROCK-FILTER-SENSITIVE-INFO: !Ref pBedrockSensitiveInfoFilterParams SRA-BEDROCK-CENTRAL-OBSERVABILITY: !Ref pBedrockCentralObservabilityParams + SRA-BEDROCK-CHECK-KB-LOGGING: !Ref pBedrockKBLoggingRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From 411f4801fc83fda59c6e16147c5dba61e4464d50 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Feb 2025 16:05:28 -0700 Subject: [PATCH 02/36] fix config rule --- .../lambda/rules/sra_bedrock_check_kb_logging/app.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index 53e67c11e..52f5b26c8 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -25,7 +25,7 @@ AWS_REGION = os.environ.get("AWS_REGION") # Initialize AWS clients -bedrock_client = boto3.client("bedrock", region_name=AWS_REGION) +bedrock_agent_client = boto3.client("bedrock-agent", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) @@ -41,9 +41,9 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: try: # List all knowledge bases kb_list = [] - paginator = bedrock_client.get_paginator('list_knowledge_bases') + paginator = bedrock_agent_client.get_paginator('list_knowledge_bases') for page in paginator.paginate(): - kb_list.extend(page.get('knowledgeBases', [])) + kb_list.extend(page.get('knowledgeBaseSummaries', [])) if not kb_list: return "COMPLIANT", "No knowledge bases found in the account" @@ -54,7 +54,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: for kb in kb_list: kb_id = kb['knowledgeBaseId'] try: - kb_details = bedrock_client.get_knowledge_base( + kb_details = bedrock_agent_client.get_knowledge_base( knowledgeBaseId=kb_id ) From de6e5de7b0d2eb60b6f25fb202f05166d24217f9 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Feb 2025 19:56:52 -0700 Subject: [PATCH 03/36] adding kb ingestion encryption check rule --- .../app.py | 108 ++++++++++++++++++ .../genai/bedrock_org/lambda/src/app.py | 2 + .../sra_config_lambda_iam_permissions.json | 16 +++ .../templates/sra-bedrock-org-main.yaml | 15 +++ 4 files changed, 141 insertions(+) create mode 100644 aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py new file mode 100644 index 000000000..6515ced01 --- /dev/null +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py @@ -0,0 +1,108 @@ +"""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 evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + """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: + # List all knowledge bases + 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) + + # Get data sources for each knowledge base + try: + data_sources = bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb_id + ) + + # Check if any data source is not encrypted + unencrypted_sources = [] + for source in data_sources.get("dataSourceSummaries", []): + if not source.get("serverSideEncryptionConfiguration", {}).get("kmsKeyArn"): + unencrypted_sources.append(source.get("name", source["dataSourceId"])) + + if unencrypted_sources: + non_compliant_kbs.append(f"{kb_name} (unencrypted sources: {', '.join(unencrypted_sources)})") + + except ClientError as e: + LOGGER.error(f"Error checking data sources for knowledge base {kb_name}: {str(e)}") + if e.response["Error"]["Code"] == "AccessDeniedException": + non_compliant_kbs.append(f"{kb_name} (access denied)") + else: + raise + + if non_compliant_kbs: + return "NON_COMPLIANT", f"The following knowledge bases have unencrypted data sources: {'; '.join(non_compliant_kbs)}" + return "COMPLIANT", "All knowledge base data sources are encrypted with KMS" + + except Exception as e: + LOGGER.error(f"Error evaluating Bedrock Knowledge Base encryption: {str(e)}") + return "ERROR", f"Error evaluating compliance: {str(e)}" + +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.") \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index db2ff6c25..0a83fc7ab 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -202,6 +202,8 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\]\}$', "SRA-BEDROCK-CHECK-KB-LOGGING": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', + "SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', } # Instantiate sra class objects diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 8d829ae61..38a6f0162 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -137,5 +137,21 @@ "Resource": "*" } ] + }, + "sra-bedrock-check-kb-ingestion-encryption": { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "AllowKnowledgeBaseAccess", + "Effect": "Allow", + "Action": [ + "bedrock:ListKnowledgeBases", + "bedrock:GetKnowledgeBase", + "bedrock:ListDataSources", + "bedrock:GetDataSource" + ], + "Resource": "*" + } + ] } } diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index 2ea8fc31a..96efc0d77 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -292,6 +292,17 @@ Parameters: Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + pBedrockKBIngestionEncryptionRuleParams: + Type: String + Default: '{"deploy": "true", "accounts": ["444455556666"], "regions": ["us-west-2"], "input_params": {}}' + Description: Bedrock Knowledge Base Data Ingestion Encryption Config Rule Parameters + AllowedPattern: ^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$ + ConstraintDescription: + "Must be a valid JSON string containing: 'deploy' (true/false), 'accounts' (array of account numbers), + 'regions' (array of region names), and 'input_params' object/dict (input params must be empty). Arrays can be empty. + Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or + {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + Metadata: AWS::CloudFormation::Interface: ParameterGroups: @@ -333,6 +344,7 @@ Metadata: - pBedrockS3EndpointsRuleParams - pBedrockGuardrailEncryptionRuleParams - pBedrockKBLoggingRuleParams + - pBedrockKBIngestionEncryptionRuleParams - Label: default: Bedrock CloudWatch Metric Filters Parameters: @@ -402,6 +414,8 @@ Metadata: default: Bedrock Regions pBedrockKBLoggingRuleParams: default: Bedrock Knowledge Base Logging Config Rule Parameters + pBedrockKBIngestionEncryptionRuleParams: + default: Bedrock Knowledge Base Data Ingestion Encryption Config Rule Parameters Resources: rBedrockOrgLambdaRole: @@ -683,6 +697,7 @@ Resources: SRA-BEDROCK-FILTER-SENSITIVE-INFO: !Ref pBedrockSensitiveInfoFilterParams SRA-BEDROCK-CENTRAL-OBSERVABILITY: !Ref pBedrockCentralObservabilityParams SRA-BEDROCK-CHECK-KB-LOGGING: !Ref pBedrockKBLoggingRuleParams + SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION: !Ref pBedrockKBIngestionEncryptionRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From 742f13af05d82aaf07d21af2670479f2a9c823c7 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Sat, 1 Mar 2025 12:09:36 -0700 Subject: [PATCH 04/36] added kb s3 bucket check rule --- .../sra_bedrock_check_kb_s3_bucket/app.py | 139 ++++++++++++++++++ .../genai/bedrock_org/lambda/src/app.py | 4 + .../sra_config_lambda_iam_permissions.json | 28 ++++ .../templates/sra-bedrock-org-main.yaml | 18 +++ 4 files changed, 189 insertions(+) create mode 100644 aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py new file mode 100644 index 000000000..01a10a618 --- /dev/null +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -0,0 +1,139 @@ +"""Config rule to check knowledge base S3 bucket configuration 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) +s3_client = boto3.client("s3", region_name=AWS_REGION) +config_client = boto3.client("config", region_name=AWS_REGION) + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + """Evaluate if Bedrock Knowledge Base S3 bucket has required configurations. + + Args: + rule_parameters (dict): Rule parameters from AWS Config rule. + + Returns: + tuple[str, str]: Compliance status and annotation + """ + try: + # Get all knowledge bases + non_compliant_buckets = [] + paginator = bedrock_agent_client.get_paginator("list_knowledge_bases") + + for page in paginator.paginate(): + for kb in page["knowledgeBaseSummaries"]: + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb["knowledgeBaseId"]) + data_source = bedrock_agent_client.get_data_source( + knowledgeBaseId=kb["knowledgeBaseId"], + dataSourceId=kb_details["dataSource"]["dataSourceId"] + ) + + # Extract bucket name from S3 path + s3_path = data_source["configuration"]["s3Configuration"]["bucketName"] + bucket_name = s3_path.split("/")[0] + + issues = [] + + # Check retention + if rule_parameters.get("check_retention", "true").lower() == "true": + try: + lifecycle = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name) + if not any(rule.get("Expiration") for rule in lifecycle.get("Rules", [])): + issues.append("retention") + except ClientError as e: + if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration": + issues.append("retention") + + # Check encryption + if rule_parameters.get("check_encryption", "true").lower() == "true": + try: + encryption = s3_client.get_bucket_encryption(Bucket=bucket_name) + if not encryption.get("ServerSideEncryptionConfiguration"): + issues.append("encryption") + except ClientError: + issues.append("encryption") + + # Check server access logging + if rule_parameters.get("check_access_logging", "true").lower() == "true": + logging_config = s3_client.get_bucket_logging(Bucket=bucket_name) + if not logging_config.get("LoggingEnabled"): + issues.append("access logging") + + # Check object lock + if rule_parameters.get("check_object_locking", "true").lower() == "true": + try: + lock_config = s3_client.get_bucket_object_lock_configuration(Bucket=bucket_name) + if not lock_config.get("ObjectLockConfiguration"): + issues.append("object locking") + except ClientError: + issues.append("object locking") + + # Check versioning + if rule_parameters.get("check_versioning", "true").lower() == "true": + versioning = s3_client.get_bucket_versioning(Bucket=bucket_name) + if versioning.get("Status") != "Enabled": + issues.append("versioning") + + if issues: + non_compliant_buckets.append(f"{bucket_name} (missing: {', '.join(issues)})") + + 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" + + except Exception as e: + LOGGER.error(f"Error evaluating Knowledge Base S3 bucket configurations: {str(e)}") + return "ERROR", f"Error evaluating compliance: {str(e)}" + +def lambda_handler(event: dict, context: Any) -> None: + """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"]) + + LOGGER.info("Compliance evaluation complete.") \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 0a83fc7ab..9b561e863 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -204,6 +204,10 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', "SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', + "SRA-BEDROCK-CHECK-KB-S3-BUCKET": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*\{(\s*"check_retention"\s*:\s*"(true|false)")?(\s*,\s*"check_encryption"\s*:\s*' + + r'"(true|false)")?(\s*,\s*"check_access_logging"\s*:\s*"(true|false)")?(\s*,\s*"check_object_locking"\s*:\s*"(true|false)")?(\s*,\s*' + + r'"check_versioning"\s*:\s*"(true|false)")?\s*\}\}$', } # Instantiate sra class objects diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 38a6f0162..c7a8b1856 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -153,5 +153,33 @@ "Resource": "*" } ] + }, + "sra-bedrock-check-kb-s3-bucket": { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "AllowKnowledgeBaseAccess", + "Effect": "Allow", + "Action": [ + "bedrock:ListKnowledgeBases", + "bedrock:GetKnowledgeBase", + "bedrock:ListDataSources", + "bedrock:GetDataSource" + ], + "Resource": "*" + }, + { + "Sid": "AllowS3BucketAccess", + "Effect": "Allow", + "Action": [ + "s3:GetBucketLifecycleConfiguration", + "s3:GetBucketEncryption", + "s3:GetBucketLogging", + "s3:GetBucketObjectLockConfiguration", + "s3:GetBucketVersioning" + ], + "Resource": "arn:aws:s3:::*" + } + ] } } diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index 96efc0d77..ce8984f77 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -303,6 +303,20 @@ Parameters: Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + pBedrockKBS3BucketRuleParams: + Type: String + Default: '{"deploy": "true", "accounts": ["444455556666"], "regions": ["us-west-2"], "input_params": {"check_retention": "true", "check_encryption": "true", "check_access_logging": "true", "check_object_locking": "true", "check_versioning": "true"}}' + Description: Bedrock Knowledge Base S3 Bucket Config Rule Parameters + AllowedPattern: ^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*\{(\s*"check_retention"\s*:\s*"(true|false)")?(\s*,\s*"check_encryption"\s*:\s*"(true|false)")?(\s*,\s*"check_access_logging"\s*:\s*"(true|false)")?(\s*,\s*"check_object_locking"\s*:\s*"(true|false)")?(\s*,\s*"check_versioning"\s*:\s*"(true|false)")?\s*\}\}$ + ConstraintDescription: > + Must be a valid JSON string containing: 'deploy' (true/false), 'accounts' (array of account numbers), + 'regions' (array of region names), and 'input_params' object with optional parameters: + 'check_retention', 'check_encryption', 'check_access_logging', 'check_object_locking', 'check_versioning'. + Each parameter in 'input_params' should be either "true" or "false". + Arrays can be empty. + Example: {"deploy": "true", "accounts": ["123456789012"], "regions": ["us-east-1"], "input_params": {"check_retention": "true", "check_encryption": "true", "check_access_logging": "true", "check_object_locking": "true", "check_versioning": "true"}} or + {"deploy": "false", "accounts": [], "regions": [], "input_params": {}} + Metadata: AWS::CloudFormation::Interface: ParameterGroups: @@ -345,6 +359,7 @@ Metadata: - pBedrockGuardrailEncryptionRuleParams - pBedrockKBLoggingRuleParams - pBedrockKBIngestionEncryptionRuleParams + - pBedrockKBS3BucketRuleParams - Label: default: Bedrock CloudWatch Metric Filters Parameters: @@ -416,6 +431,8 @@ Metadata: default: Bedrock Knowledge Base Logging Config Rule Parameters pBedrockKBIngestionEncryptionRuleParams: default: Bedrock Knowledge Base Data Ingestion Encryption Config Rule Parameters + pBedrockKBS3BucketRuleParams: + default: Bedrock Knowledge Base S3 Bucket Config Rule Parameters Resources: rBedrockOrgLambdaRole: @@ -698,6 +715,7 @@ Resources: SRA-BEDROCK-CENTRAL-OBSERVABILITY: !Ref pBedrockCentralObservabilityParams SRA-BEDROCK-CHECK-KB-LOGGING: !Ref pBedrockKBLoggingRuleParams SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION: !Ref pBedrockKBIngestionEncryptionRuleParams + SRA-BEDROCK-CHECK-KB-S3-BUCKET: !Ref pBedrockKBS3BucketRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From 15e48186cb56e8beb26c53052af3e283ad426f41 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 4 Mar 2025 15:35:51 -0700 Subject: [PATCH 05/36] added prototype config rule for kb vector store secrete --- .../app.py | 116 ++++++++++++++++++ .../genai/bedrock_org/lambda/src/app.py | 2 + .../sra_config_lambda_iam_permissions.json | 22 ++++ .../templates/sra-bedrock-org-main.yaml | 15 +++ 4 files changed, 155 insertions(+) create mode 100644 aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py new file mode 100644 index 000000000..7d2c83f46 --- /dev/null +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py @@ -0,0 +1,116 @@ +"""Config rule to check knowledge base vector store secret configuration 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) +secretsmanager_client = boto3.client("secretsmanager", region_name=AWS_REGION) +config_client = boto3.client("config", region_name=AWS_REGION) + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + """Evaluate if Bedrock Knowledge Base vector stores are using KMS encrypted secrets. + + 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) + + try: + # Get knowledge base details + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) + vector_store = kb_details.get("vectorStoreConfiguration") + + if vector_store: + secret_arn = vector_store.get("secretArn") + if not secret_arn: + non_compliant_kbs.append(f"{kb_name} (no secret configured)") + continue + + try: + # Check if secret uses CMK + secret_details = secretsmanager_client.describe_secret(SecretId=secret_arn) + if not secret_details.get("KmsKeyId"): + non_compliant_kbs.append(f"{kb_name} (secret not using CMK)") + except ClientError as e: + LOGGER.error(f"Error checking secret {secret_arn}: {str(e)}") + if e.response["Error"]["Code"] == "AccessDeniedException": + non_compliant_kbs.append(f"{kb_name} (secret access denied)") + else: + raise + + except ClientError as e: + LOGGER.error(f"Error checking knowledge base {kb_name}: {str(e)}") + if e.response["Error"]["Code"] == "AccessDeniedException": + non_compliant_kbs.append(f"{kb_name} (access denied)") + else: + raise + + if non_compliant_kbs: + return "NON_COMPLIANT", f"The following knowledge bases have vector store secret issues: {'; '.join(non_compliant_kbs)}" + return "COMPLIANT", "All knowledge base vector stores are using KMS encrypted secrets" + + except Exception as e: + LOGGER.error(f"Error evaluating Bedrock Knowledge Base vector store secrets: {str(e)}") + return "ERROR", f"Error evaluating compliance: {str(e)}" + +def lambda_handler(event: dict, context: Any) -> None: + """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"]) + + LOGGER.info("Compliance evaluation complete.") \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 9b561e863..42b1b5fed 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -208,6 +208,8 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*\{(\s*"check_retention"\s*:\s*"(true|false)")?(\s*,\s*"check_encryption"\s*:\s*' + r'"(true|false)")?(\s*,\s*"check_access_logging"\s*:\s*"(true|false)")?(\s*,\s*"check_object_locking"\s*:\s*"(true|false)")?(\s*,\s*' + r'"check_versioning"\s*:\s*"(true|false)")?\s*\}\}$', + "SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*' + + r'"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', } # Instantiate sra class objects diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index c7a8b1856..f7cf42ff0 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -181,5 +181,27 @@ "Resource": "arn:aws:s3:::*" } ] + }, + "sra-bedrock-check-kb-vector-store-secret": { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "AllowKnowledgeBaseAccess", + "Effect": "Allow", + "Action": [ + "bedrock:ListKnowledgeBases", + "bedrock:GetKnowledgeBase" + ], + "Resource": "*" + }, + { + "Sid": "AllowSecretsManagerAccess", + "Effect": "Allow", + "Action": [ + "secretsmanager:DescribeSecret" + ], + "Resource": "*" + } + ] } } diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index ce8984f77..97a7e0fcd 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -317,6 +317,17 @@ Parameters: Example: {"deploy": "true", "accounts": ["123456789012"], "regions": ["us-east-1"], "input_params": {"check_retention": "true", "check_encryption": "true", "check_access_logging": "true", "check_object_locking": "true", "check_versioning": "true"}} or {"deploy": "false", "accounts": [], "regions": [], "input_params": {}} + pBedrockKBVectorStoreSecretRuleParams: + Type: String + Default: '{"deploy": "true", "accounts": ["444455556666"], "regions": ["us-west-2"], "input_params": {}}' + Description: Bedrock Knowledge Base Vector Store Secret Config Rule Parameters + AllowedPattern: ^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$ + ConstraintDescription: + "Must be a valid JSON string containing: 'deploy' (true/false), 'accounts' (array of account numbers), + 'regions' (array of region names), and 'input_params' object/dict (input params must be empty). Arrays can be empty. + Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or + {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + Metadata: AWS::CloudFormation::Interface: ParameterGroups: @@ -360,6 +371,7 @@ Metadata: - pBedrockKBLoggingRuleParams - pBedrockKBIngestionEncryptionRuleParams - pBedrockKBS3BucketRuleParams + - pBedrockKBVectorStoreSecretRuleParams - Label: default: Bedrock CloudWatch Metric Filters Parameters: @@ -433,6 +445,8 @@ Metadata: default: Bedrock Knowledge Base Data Ingestion Encryption Config Rule Parameters pBedrockKBS3BucketRuleParams: default: Bedrock Knowledge Base S3 Bucket Config Rule Parameters + pBedrockKBVectorStoreSecretRuleParams: + default: Bedrock Knowledge Base Vector Store Secret Config Rule Parameters Resources: rBedrockOrgLambdaRole: @@ -716,6 +730,7 @@ Resources: SRA-BEDROCK-CHECK-KB-LOGGING: !Ref pBedrockKBLoggingRuleParams SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION: !Ref pBedrockKBIngestionEncryptionRuleParams SRA-BEDROCK-CHECK-KB-S3-BUCKET: !Ref pBedrockKBS3BucketRuleParams + SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET: !Ref pBedrockKBVectorStoreSecretRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From 652123cc44c3fc7dcb314767efc60de205ad7d14 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Sat, 8 Mar 2025 17:41:18 -0700 Subject: [PATCH 06/36] added check for opensearch encryption (if used in KB) --- .../app.py | 150 ++++++++++++++++++ .../genai/bedrock_org/lambda/src/app.py | 1 + .../sra_config_lambda_iam_permissions.json | 23 +++ .../templates/sra-bedrock-org-main.yaml | 15 ++ 4 files changed, 189 insertions(+) create mode 100644 aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py new file mode 100644 index 000000000..bb91ad81c --- /dev/null +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -0,0 +1,150 @@ +"""Config rule to check OpenSearch vector store encryption for Bedrock Knowledge Base. + +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) +opensearch_client = boto3.client("opensearch", region_name=AWS_REGION) +config_client = boto3.client("config", region_name=AWS_REGION) + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + """Evaluate if Bedrock Knowledge Base OpenSearch vector stores are encrypted with KMS CMK. + + 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) + + try: + # Get knowledge base details + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) + vector_store = kb_details.get("vectorStoreConfiguration") + + if vector_store and vector_store.get("vectorStoreType") == "OPENSEARCH": + # Extract OpenSearch domain information + opensearch_config = vector_store.get("opensearchServerlessConfiguration") or vector_store.get("opensearchConfiguration") + + if not opensearch_config: + non_compliant_kbs.append(f"{kb_name} (missing OpenSearch configuration)") + continue + + # Check if it's OpenSearch Serverless or standard OpenSearch + if "collectionArn" in opensearch_config: + # OpenSearch Serverless - always encrypted with AWS owned key at minimum + collection_id = opensearch_config["collectionArn"].split("/")[-1] + try: + collection = opensearch_client.get_security_policy( + Name=collection_id, + Type="encryption" + ) + # Check if using customer managed key + security_policy = collection.get("securityPolicyDetail", {}) + if security_policy.get("Type") == "encryption": + encryption_policy = security_policy.get("SecurityPolicies", [])[0] + kms_key_arn = encryption_policy.get("KmsARN", "") + + # If not using customer managed key + if not kms_key_arn or "aws/opensearchserverless" in kms_key_arn: + non_compliant_kbs.append(f"{kb_name} (OpenSearch Serverless not using CMK)") + except ClientError as e: + LOGGER.error(f"Error checking OpenSearch Serverless collection: {str(e)}") + non_compliant_kbs.append(f"{kb_name} (error checking OpenSearch Serverless)") + else: + # Standard OpenSearch + domain_endpoint = opensearch_config.get("endpoint", "") + if not domain_endpoint: + non_compliant_kbs.append(f"{kb_name} (missing OpenSearch domain endpoint)") + continue + + # Extract domain name from endpoint + domain_name = domain_endpoint.split(".")[0] + + try: + domain = opensearch_client.describe_domain(DomainName=domain_name) + encryption_config = domain.get("DomainStatus", {}).get("EncryptionAtRestOptions", {}) + + # Check if encryption is enabled and using CMK + if not encryption_config.get("Enabled", False): + non_compliant_kbs.append(f"{kb_name} (OpenSearch domain encryption not enabled)") + elif not encryption_config.get("KmsKeyId"): + non_compliant_kbs.append(f"{kb_name} (OpenSearch domain not using CMK)") + except ClientError as e: + LOGGER.error(f"Error checking OpenSearch domain: {str(e)}") + non_compliant_kbs.append(f"{kb_name} (error checking OpenSearch domain)") + + except ClientError as e: + LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") + if e.response["Error"]["Code"] == "AccessDeniedException": + non_compliant_kbs.append(f"{kb_name} (access denied)") + else: + raise + + if non_compliant_kbs: + return "NON_COMPLIANT", f"The following knowledge bases have OpenSearch vector stores not encrypted with CMK: {'; '.join(non_compliant_kbs)}" + return "COMPLIANT", "All knowledge base OpenSearch vector stores are encrypted with KMS CMK" + + except Exception as e: + LOGGER.error(f"Error evaluating Bedrock Knowledge Base OpenSearch encryption: {str(e)}") + return "ERROR", f"Error evaluating compliance: {str(e)}" + +def lambda_handler(event: dict, context: Any) -> None: + """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"]) + + LOGGER.info("Compliance evaluation complete.") \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 42b1b5fed..7261aa9fd 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -210,6 +210,7 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'"check_versioning"\s*:\s*"(true|false)")?\s*\}\}$', "SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*' + r'"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', + "SRA-BEDROCK-CHECK-KB-OPENSEARCH-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', } # Instantiate sra class objects diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index f7cf42ff0..b7fefff75 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -203,5 +203,28 @@ "Resource": "*" } ] + }, + "sra-bedrock-check-kb-opensearch-encryption": { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "AllowKnowledgeBaseAccess", + "Effect": "Allow", + "Action": [ + "bedrock:ListKnowledgeBases", + "bedrock:GetKnowledgeBase" + ], + "Resource": "*" + }, + { + "Sid": "AllowOpenSearchAccess", + "Effect": "Allow", + "Action": [ + "opensearch:DescribeDomain", + "opensearch:GetSecurityPolicy" + ], + "Resource": "*" + } + ] } } diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index 97a7e0fcd..9e0ce3d9a 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -328,6 +328,17 @@ Parameters: Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + pBedrockKBOpenSearchEncryptionRuleParams: + Type: String + Default: '{"deploy": "true", "accounts": ["444455556666"], "regions": ["us-west-2"], "input_params": {}}' + Description: Bedrock Knowledge Base OpenSearch Encryption Config Rule Parameters + AllowedPattern: ^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$ + ConstraintDescription: + "Must be a valid JSON string containing: 'deploy' (true/false), 'accounts' (array of account numbers), + 'regions' (array of region names), and 'input_params' object/dict (input params must be empty). Arrays can be empty. + Example: {\"deploy\": \"true\", \"accounts\": [\"123456789012\"], \"regions\": [\"us-east-1\"], \"input_params\": {}} or + {\"deploy\": \"false\", \"accounts\": [], \"regions\": [], \"input_params\": {}}" + Metadata: AWS::CloudFormation::Interface: ParameterGroups: @@ -372,6 +383,7 @@ Metadata: - pBedrockKBIngestionEncryptionRuleParams - pBedrockKBS3BucketRuleParams - pBedrockKBVectorStoreSecretRuleParams + - pBedrockKBOpenSearchEncryptionRuleParams - Label: default: Bedrock CloudWatch Metric Filters Parameters: @@ -447,6 +459,8 @@ Metadata: default: Bedrock Knowledge Base S3 Bucket Config Rule Parameters pBedrockKBVectorStoreSecretRuleParams: default: Bedrock Knowledge Base Vector Store Secret Config Rule Parameters + pBedrockKBOpenSearchEncryptionRuleParams: + default: Bedrock Knowledge Base OpenSearch Encryption Config Rule Parameters Resources: rBedrockOrgLambdaRole: @@ -731,6 +745,7 @@ Resources: SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION: !Ref pBedrockKBIngestionEncryptionRuleParams SRA-BEDROCK-CHECK-KB-S3-BUCKET: !Ref pBedrockKBS3BucketRuleParams SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET: !Ref pBedrockKBVectorStoreSecretRuleParams + SRA-BEDROCK-CHECK-KB-OPEN-SEARCH-ENCRYPTION: !Ref pBedrockKBOpenSearchEncryptionRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From 089feeb4a92683dc235e6f62abcc77cc154d4033 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 11 Mar 2025 09:32:42 -0600 Subject: [PATCH 07/36] update param name --- .../genai/bedrock_org/templates/sra-bedrock-org-main.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml index 9e0ce3d9a..33417d041 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml +++ b/aws_sra_examples/solutions/genai/bedrock_org/templates/sra-bedrock-org-main.yaml @@ -745,7 +745,7 @@ Resources: SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION: !Ref pBedrockKBIngestionEncryptionRuleParams SRA-BEDROCK-CHECK-KB-S3-BUCKET: !Ref pBedrockKBS3BucketRuleParams SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET: !Ref pBedrockKBVectorStoreSecretRuleParams - SRA-BEDROCK-CHECK-KB-OPEN-SEARCH-ENCRYPTION: !Ref pBedrockKBOpenSearchEncryptionRuleParams + SRA-BEDROCK-CHECK-KB-OPENSEARCH-ENCRYPTION: !Ref pBedrockKBOpenSearchEncryptionRuleParams rBedrockOrgLambdaInvokePermission: Type: AWS::Lambda::Permission From e3335c2be4f16e7d5f244999849494c3bbce9ca9 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 11 Mar 2025 22:58:02 -0600 Subject: [PATCH 08/36] working to fix bug --- .../sra_bedrock_check_kb_s3_bucket/app.py | 139 +++++++++++------- 1 file changed, 86 insertions(+), 53 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index 01a10a618..654e30d97 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -45,60 +45,92 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: for page in paginator.paginate(): for kb in page["knowledgeBaseSummaries"]: - kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb["knowledgeBaseId"]) - data_source = bedrock_agent_client.get_data_source( - knowledgeBaseId=kb["knowledgeBaseId"], - dataSourceId=kb_details["dataSource"]["dataSourceId"] - ) + kb_id = kb["knowledgeBaseId"] - # Extract bucket name from S3 path - s3_path = data_source["configuration"]["s3Configuration"]["bucketName"] - bucket_name = s3_path.split("/")[0] + # List data sources for this knowledge base + data_sources_paginator = bedrock_agent_client.get_paginator("list_data_sources") - issues = [] - - # Check retention - if rule_parameters.get("check_retention", "true").lower() == "true": - try: - lifecycle = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name) - if not any(rule.get("Expiration") for rule in lifecycle.get("Rules", [])): - issues.append("retention") - except ClientError as e: - if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration": - issues.append("retention") - - # Check encryption - if rule_parameters.get("check_encryption", "true").lower() == "true": - try: - encryption = s3_client.get_bucket_encryption(Bucket=bucket_name) - if not encryption.get("ServerSideEncryptionConfiguration"): - issues.append("encryption") - except ClientError: - issues.append("encryption") - - # Check server access logging - if rule_parameters.get("check_access_logging", "true").lower() == "true": - logging_config = s3_client.get_bucket_logging(Bucket=bucket_name) - if not logging_config.get("LoggingEnabled"): - issues.append("access logging") - - # Check object lock - if rule_parameters.get("check_object_locking", "true").lower() == "true": - try: - lock_config = s3_client.get_bucket_object_lock_configuration(Bucket=bucket_name) - if not lock_config.get("ObjectLockConfiguration"): - issues.append("object locking") - except ClientError: - issues.append("object locking") - - # Check versioning - if rule_parameters.get("check_versioning", "true").lower() == "true": - versioning = s3_client.get_bucket_versioning(Bucket=bucket_name) - if versioning.get("Status") != "Enabled": - issues.append("versioning") - - if issues: - non_compliant_buckets.append(f"{bucket_name} (missing: {', '.join(issues)})") + for ds_page in data_sources_paginator.paginate(knowledgeBaseId=kb_id): + for ds in ds_page.get("dataSourceSummaries", []): + data_source = bedrock_agent_client.get_data_source( + knowledgeBaseId=kb_id, + dataSourceId=ds["dataSourceId"] + ) + + # Check if this is an S3 data source and extract bucket name + LOGGER.info(f"Data source structure: {json.dumps(data_source)}") + if "s3Configuration" in data_source.get("dataSource", {}).get("dataSourceConfiguration", {}): + s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] + bucket_name = s3_config.get("bucketName", "") + else: + continue + + if not bucket_name: + LOGGER.info(f"No bucket name found for data source {ds['dataSourceId']}") + continue + + # If bucket name contains a path, extract just the bucket name + if "/" in bucket_name: + bucket_name = bucket_name.split("/")[0] + + LOGGER.info(f"Checking S3 bucket: {bucket_name}") + + issues = [] + + # Check retention + if rule_parameters.get("check_retention", "true").lower() == "true": + try: + lifecycle = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name) + if not any(rule.get("Expiration") for rule in lifecycle.get("Rules", [])): + issues.append("retention") + except ClientError as e: + if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration": + issues.append("retention") + elif e.response["Error"]["Code"] != "NoSuchBucket": + LOGGER.error(f"Error checking retention for bucket {bucket_name}: {str(e)}") + + # Check encryption + if rule_parameters.get("check_encryption", "true").lower() == "true": + try: + encryption = s3_client.get_bucket_encryption(Bucket=bucket_name) + if not encryption.get("ServerSideEncryptionConfiguration"): + issues.append("encryption") + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + issues.append("encryption") + + # Check server access logging + if rule_parameters.get("check_access_logging", "true").lower() == "true": + try: + logging_config = s3_client.get_bucket_logging(Bucket=bucket_name) + if not logging_config.get("LoggingEnabled"): + issues.append("access logging") + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + issues.append("access logging") + + # Check object lock + if rule_parameters.get("check_object_locking", "true").lower() == "true": + try: + lock_config = s3_client.get_object_lock_configuration(Bucket=bucket_name) + if not lock_config.get("ObjectLockConfiguration"): + issues.append("object locking") + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + issues.append("object locking") + + # Check versioning + if rule_parameters.get("check_versioning", "true").lower() == "true": + try: + versioning = s3_client.get_bucket_versioning(Bucket=bucket_name) + if versioning.get("Status") != "Enabled": + issues.append("versioning") + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + issues.append("versioning") + + if issues: + non_compliant_buckets.append(f"{bucket_name} (missing: {', '.join(issues)})") if non_compliant_buckets: return "NON_COMPLIANT", f"The following KB S3 buckets are non-compliant: {'; '.join(non_compliant_buckets)}" @@ -136,4 +168,5 @@ def lambda_handler(event: dict, context: Any) -> None: config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) - LOGGER.info("Compliance evaluation complete.") \ No newline at end of file + LOGGER.info("Compliance evaluation complete.") + \ No newline at end of file From d9c29f0f0abcfdc7d04f8200e7672fd80805bf44 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 13 Mar 2025 11:27:03 -0600 Subject: [PATCH 09/36] updating lambda code. updating permissions --- .../sra_bedrock_check_kb_s3_bucket/app.py | 44 ++++++++++++++----- .../sra_config_lambda_iam_permissions.json | 6 +-- 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index 654e30d97..03148bc87 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -11,6 +11,7 @@ import logging import os from typing import Any +from datetime import datetime import boto3 from botocore.exceptions import ClientError @@ -58,21 +59,40 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: ) # Check if this is an S3 data source and extract bucket name - LOGGER.info(f"Data source structure: {json.dumps(data_source)}") - if "s3Configuration" in data_source.get("dataSource", {}).get("dataSourceConfiguration", {}): - s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] - bucket_name = s3_config.get("bucketName", "") - else: - continue + try: + # Use a custom JSON encoder to handle datetime objects + class DateTimeEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, datetime): + return obj.isoformat() + return super().default(obj) - if not bucket_name: - LOGGER.info(f"No bucket name found for data source {ds['dataSourceId']}") + LOGGER.info(f"Data source structure: {json.dumps(data_source, cls=DateTimeEncoder)}") + + if "dataSource" in data_source and "dataSourceConfiguration" in data_source["dataSource"]: + if "s3Configuration" in data_source["dataSource"]["dataSourceConfiguration"]: + s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] + bucket_arn = s3_config.get("bucketArn", "") + + if not bucket_arn: + LOGGER.info(f"No bucket ARN found for data source {ds['dataSourceId']}") + continue + + # Extract bucket name from ARN + # ARN format: arn:aws:s3:::bucket-name + bucket_name = bucket_arn.split(":")[-1] + + # If bucket name contains a path, extract just the bucket name + if "/" in bucket_name: + bucket_name = bucket_name.split("/")[0] + else: + continue + else: + continue + except Exception as e: + LOGGER.error(f"Error processing data source: {str(e)}") continue - # If bucket name contains a path, extract just the bucket name - if "/" in bucket_name: - bucket_name = bucket_name.split("/")[0] - LOGGER.info(f"Checking S3 bucket: {bucket_name}") issues = [] diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index b7fefff75..0e169c31b 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -172,10 +172,10 @@ "Sid": "AllowS3BucketAccess", "Effect": "Allow", "Action": [ - "s3:GetBucketLifecycleConfiguration", - "s3:GetBucketEncryption", - "s3:GetBucketLogging", "s3:GetBucketObjectLockConfiguration", + "s3:GetLifecycleConfiguration", + "s3:GetEncryptionConfiguration", + "s3:GetBucketLogging", "s3:GetBucketVersioning" ], "Resource": "arn:aws:s3:::*" From df55394bdcd935576e00d74685ecfc18f5eebbc8 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 20 Mar 2025 09:33:01 -0600 Subject: [PATCH 10/36] return arn if already exists --- .../solutions/genai/bedrock_org/lambda/src/sra_iam.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py index e58d2349d..18198ca87 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py @@ -113,6 +113,8 @@ def create_role(self, role_name: str, trust_policy: dict, solution_name: str) -> except ClientError as error: if error.response["Error"]["Code"] == "EntityAlreadyExists": self.LOGGER.info(f"{role_name} role already exists!") + response = self.IAM_CLIENT.get_role(RoleName=role_name) + return {"Role": {"Arn": response["Role"]["Arn"]}} return {"Role": {"Arn": "error"}} def create_policy(self, policy_name: str, policy_document: dict, solution_name: str) -> dict: From 530369ecddeef9b1d48316e1d5510b2d8d8524da Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Mon, 24 Mar 2025 15:39:53 -0600 Subject: [PATCH 11/36] update README.md file --- .../solutions/genai/bedrock_org/README.md | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/README.md b/aws_sra_examples/solutions/genai/bedrock_org/README.md index cfcf25eef..ba6356f81 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/README.md +++ b/aws_sra_examples/solutions/genai/bedrock_org/README.md @@ -102,6 +102,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 +144,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 +165,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 +386,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) From 1a42368698b2d3eaaee60a5dbfdaf0c825885c59 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Mon, 24 Mar 2025 15:58:05 -0600 Subject: [PATCH 12/36] update readme --- .../solutions/genai/bedrock_org/README.md | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/README.md b/aws_sra_examples/solutions/genai/bedrock_org/README.md index ba6356f81..80b005bc7 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/README.md +++ b/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) --- @@ -460,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 From d0b6a2a916b21d89a28e7cfa64987719331e3ac6 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 08:13:30 -0600 Subject: [PATCH 13/36] updating for flake8 and mypy errors --- .../app.py | 69 ++-- .../rules/sra_bedrock_check_kb_logging/app.py | 15 +- .../app.py | 181 ++++++----- .../sra_bedrock_check_kb_s3_bucket/app.py | 296 +++++++++++------- .../app.py | 84 ++--- 5 files changed, 395 insertions(+), 250 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 6515ced01..065a434fe 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py @@ -28,7 +28,43 @@ bedrock_agent_client = boto3.client("bedrock-agent", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) -def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + +def check_data_sources(kb_id: str, kb_name: str) -> str | None: # noqa: CFQ004 + """Check if a knowledge base's data sources are encrypted. + + 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) + if not isinstance(data_sources, dict): + return f"{kb_name} (invalid data sources response)" + unencrypted_sources = [] + for source in data_sources.get("dataSourceSummaries", []): + if not isinstance(source, dict): + continue + encryption_config = source.get("serverSideEncryptionConfiguration", {}) + if not isinstance(encryption_config, dict) or not encryption_config.get("kmsKeyArn"): + unencrypted_sources.append(source.get("name", source["dataSourceId"])) + + if unencrypted_sources: + return f"{kb_name} (unencrypted sources: {', '.join(unencrypted_sources)})" + 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: @@ -38,36 +74,16 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: tuple[str, str]: Compliance type and annotation message. """ try: - # List all knowledge bases 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) - - # Get data sources for each knowledge base - try: - data_sources = bedrock_agent_client.list_data_sources( - knowledgeBaseId=kb_id - ) - - # Check if any data source is not encrypted - unencrypted_sources = [] - for source in data_sources.get("dataSourceSummaries", []): - if not source.get("serverSideEncryptionConfiguration", {}).get("kmsKeyArn"): - unencrypted_sources.append(source.get("name", source["dataSourceId"])) - - if unencrypted_sources: - non_compliant_kbs.append(f"{kb_name} (unencrypted sources: {', '.join(unencrypted_sources)})") - - except ClientError as e: - LOGGER.error(f"Error checking data sources for knowledge base {kb_name}: {str(e)}") - if e.response["Error"]["Code"] == "AccessDeniedException": - non_compliant_kbs.append(f"{kb_name} (access denied)") - else: - raise + error = check_data_sources(kb_id, kb_name) + if error: + non_compliant_kbs.append(error) if non_compliant_kbs: return "NON_COMPLIANT", f"The following knowledge bases have unencrypted data sources: {'; '.join(non_compliant_kbs)}" @@ -77,6 +93,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: LOGGER.error(f"Error evaluating Bedrock Knowledge Base encryption: {str(e)}") return "ERROR", f"Error evaluating compliance: {str(e)}" + def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 """Lambda handler. @@ -105,4 +122,4 @@ def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) # type: ignore - LOGGER.info("Compliance evaluation complete.") \ No newline at end of file + LOGGER.info("Compliance evaluation complete.") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index 52f5b26c8..af8fcf725 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -29,12 +29,15 @@ config_client = boto3.client("config", region_name=AWS_REGION) -def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ004, U100 """Evaluate if Bedrock Knowledge Base logging is properly configured. Args: rule_parameters (dict): Rule parameters from AWS Config rule. + Raises: + ClientError: If there is an error checking the knowledge base + Returns: tuple[str, str]: Compliance type and annotation message. """ @@ -49,7 +52,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: return "COMPLIANT", "No knowledge bases found in the account" non_compliant_kbs = [] - + # Check each knowledge base for logging configuration for kb in kb_list: kb_id = kb['knowledgeBaseId'] @@ -57,12 +60,12 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: kb_details = bedrock_agent_client.get_knowledge_base( knowledgeBaseId=kb_id ) - + # Check if logging is enabled logging_config = kb_details.get('loggingConfiguration', {}) - if not logging_config or not logging_config.get('enabled', False): + if not isinstance(logging_config, dict) or not logging_config.get('enabled', False): non_compliant_kbs.append(f"{kb_id} ({kb.get('name', 'unnamed')})") - + except ClientError as e: LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") if e.response['Error']['Code'] == 'AccessDeniedException': @@ -107,4 +110,4 @@ def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) # type: ignore - LOGGER.info("Compliance evaluation complete.") \ No newline at end of file + LOGGER.info("Compliance evaluation complete.") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index bb91ad81c..76018167b 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -27,9 +27,107 @@ # Initialize AWS clients bedrock_agent_client = boto3.client("bedrock-agent", region_name=AWS_REGION) opensearch_client = boto3.client("opensearch", region_name=AWS_REGION) +opensearch_serverless_client = boto3.client("opensearchserverless", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) -def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + +def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: + """Check OpenSearch Serverless collection encryption. + + Args: + collection_id (str): Collection ID + kb_name (str): Knowledge base name + + Returns: + str | None: Error message if non-compliant, None if compliant + """ + try: + collection = opensearch_serverless_client.get_security_policy( + name=collection_id, + type="encryption" + ) + security_policy = collection.get("securityPolicyDetail", {}) + if security_policy.get("Type") == "encryption": + security_policies = security_policy.get("SecurityPolicies", []) + if isinstance(security_policies, list) and security_policies: + encryption_policy = security_policies[0] + kms_key_arn = encryption_policy.get("KmsARN", "") + if not kms_key_arn or "aws/opensearchserverless" in kms_key_arn: + return f"{kb_name} (OpenSearch Serverless not using CMK)" + except ClientError as e: + LOGGER.error(f"Error checking OpenSearch Serverless collection: {str(e)}") + return f"{kb_name} (error checking OpenSearch Serverless)" + return None + + +def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # noqa: CFQ004 + """Check standard OpenSearch domain encryption. + + Args: + domain_name (str): Domain name + kb_name (str): Knowledge base name + + Returns: + str | None: Error message if non-compliant, None if compliant + """ + try: + 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)" + if not encryption_config.get("KmsKeyId"): + 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: # noqa: CFQ004 + """Check a knowledge base's OpenSearch configuration. + + 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: + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) + vector_store = kb_details.get("vectorStoreConfiguration") + + if not vector_store or not isinstance(vector_store, dict): + return None + + if vector_store.get("vectorStoreType") != "OPENSEARCH": + return None + + opensearch_config = vector_store.get("opensearchServerlessConfiguration") or vector_store.get("opensearchConfiguration") + if not opensearch_config: + return 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) + + domain_endpoint = opensearch_config.get("endpoint", "") + if not domain_endpoint: + return f"{kb_name} (missing OpenSearch domain endpoint)" + domain_name = domain_endpoint.split(".")[0] + return 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)" + raise + + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 """Evaluate if Bedrock Knowledge Base OpenSearch vector stores are encrypted with KMS CMK. Args: @@ -41,85 +139,28 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: 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) - - try: - # Get knowledge base details - kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) - vector_store = kb_details.get("vectorStoreConfiguration") - - if vector_store and vector_store.get("vectorStoreType") == "OPENSEARCH": - # Extract OpenSearch domain information - opensearch_config = vector_store.get("opensearchServerlessConfiguration") or vector_store.get("opensearchConfiguration") - - if not opensearch_config: - non_compliant_kbs.append(f"{kb_name} (missing OpenSearch configuration)") - continue - - # Check if it's OpenSearch Serverless or standard OpenSearch - if "collectionArn" in opensearch_config: - # OpenSearch Serverless - always encrypted with AWS owned key at minimum - collection_id = opensearch_config["collectionArn"].split("/")[-1] - try: - collection = opensearch_client.get_security_policy( - Name=collection_id, - Type="encryption" - ) - # Check if using customer managed key - security_policy = collection.get("securityPolicyDetail", {}) - if security_policy.get("Type") == "encryption": - encryption_policy = security_policy.get("SecurityPolicies", [])[0] - kms_key_arn = encryption_policy.get("KmsARN", "") - - # If not using customer managed key - if not kms_key_arn or "aws/opensearchserverless" in kms_key_arn: - non_compliant_kbs.append(f"{kb_name} (OpenSearch Serverless not using CMK)") - except ClientError as e: - LOGGER.error(f"Error checking OpenSearch Serverless collection: {str(e)}") - non_compliant_kbs.append(f"{kb_name} (error checking OpenSearch Serverless)") - else: - # Standard OpenSearch - domain_endpoint = opensearch_config.get("endpoint", "") - if not domain_endpoint: - non_compliant_kbs.append(f"{kb_name} (missing OpenSearch domain endpoint)") - continue - - # Extract domain name from endpoint - domain_name = domain_endpoint.split(".")[0] - - try: - domain = opensearch_client.describe_domain(DomainName=domain_name) - encryption_config = domain.get("DomainStatus", {}).get("EncryptionAtRestOptions", {}) - - # Check if encryption is enabled and using CMK - if not encryption_config.get("Enabled", False): - non_compliant_kbs.append(f"{kb_name} (OpenSearch domain encryption not enabled)") - elif not encryption_config.get("KmsKeyId"): - non_compliant_kbs.append(f"{kb_name} (OpenSearch domain not using CMK)") - except ClientError as e: - LOGGER.error(f"Error checking OpenSearch domain: {str(e)}") - non_compliant_kbs.append(f"{kb_name} (error checking OpenSearch domain)") - - except ClientError as e: - LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") - if e.response["Error"]["Code"] == "AccessDeniedException": - non_compliant_kbs.append(f"{kb_name} (access denied)") - else: - raise + error = check_knowledge_base(kb_id, kb_name) + if error: + non_compliant_kbs.append(error) if non_compliant_kbs: - return "NON_COMPLIANT", f"The following knowledge bases have OpenSearch vector stores not encrypted with CMK: {'; '.join(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" except Exception as e: LOGGER.error(f"Error evaluating Bedrock Knowledge Base OpenSearch encryption: {str(e)}") return "ERROR", f"Error evaluating compliance: {str(e)}" -def lambda_handler(event: dict, context: Any) -> None: + +def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 """Lambda handler. Args: @@ -145,6 +186,6 @@ def lambda_handler(event: dict, context: Any) -> None: LOGGER.info(f"Compliance evaluation result: {compliance_type}") LOGGER.info(f"Annotation: {annotation}") - config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) + config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) # type: ignore - LOGGER.info("Compliance evaluation complete.") \ No newline at end of file + LOGGER.info("Compliance evaluation complete.") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index 03148bc87..e4679bbea 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -10,11 +10,10 @@ import json import logging import os -from typing import Any -from datetime import datetime - +from typing import Any, Union import boto3 from botocore.exceptions import ClientError +from mypy_boto3_bedrock_agent.type_defs import GetDataSourceResponseTypeDef # Setup Default Logger LOGGER = logging.getLogger(__name__) @@ -30,6 +29,182 @@ s3_client = boto3.client("s3", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) + +def check_retention(bucket_name: str) -> bool: + """Check if bucket has retention configuration. + + Args: + bucket_name (str): Name of the S3 bucket to check + + Returns: + bool: True if bucket has retention configuration, False otherwise + """ + try: + lifecycle = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name) + return any(rule.get("Expiration") for rule in lifecycle.get("Rules", [])) + except ClientError as e: + if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration": + return False + if e.response["Error"]["Code"] != "NoSuchBucket": + LOGGER.error(f"Error checking retention for bucket {bucket_name}: {str(e)}") + return False + + +def check_encryption(bucket_name: str) -> bool: + """Check if bucket has encryption configuration. + + Args: + bucket_name (str): Name of the S3 bucket to check + + Returns: + bool: True if bucket has encryption configuration, False otherwise + """ + try: + encryption = s3_client.get_bucket_encryption(Bucket=bucket_name) + return bool(encryption.get("ServerSideEncryptionConfiguration")) + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + return False + return False + + +def check_access_logging(bucket_name: str) -> bool: + """Check if bucket has access logging enabled. + + Args: + bucket_name (str): Name of the S3 bucket to check + + Returns: + bool: True if bucket has access logging enabled, False otherwise + """ + try: + logging_config = s3_client.get_bucket_logging(Bucket=bucket_name) + return bool(logging_config.get("LoggingEnabled")) + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + return False + return False + + +def check_object_locking(bucket_name: str) -> bool: + """Check if bucket has object locking enabled. + + Args: + bucket_name (str): Name of the S3 bucket to check + + Returns: + bool: True if bucket has object locking enabled, False otherwise + """ + try: + lock_config = s3_client.get_object_lock_configuration(Bucket=bucket_name) + return bool(lock_config.get("ObjectLockConfiguration")) + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + return False + return False + + +def check_versioning(bucket_name: str) -> bool: + """Check if bucket has versioning enabled. + + Args: + bucket_name (str): Name of the S3 bucket to check + + Returns: + bool: True if bucket has versioning enabled, False otherwise + """ + try: + versioning = s3_client.get_bucket_versioning(Bucket=bucket_name) + return versioning.get("Status") == "Enabled" + except ClientError as e: + if e.response["Error"]["Code"] != "NoSuchBucket": + return False + return False + + +def check_bucket_configuration(bucket_name: str, rule_parameters: dict) -> list[str]: + """Check S3 bucket configuration against required settings. + + Args: + bucket_name (str): Name of the S3 bucket + rule_parameters (dict): Rule parameters containing check flags + + Returns: + list[str]: List of missing configurations + """ + issues = [] + + if rule_parameters.get("check_retention", "true").lower() == "true" and not check_retention(bucket_name): + issues.append("retention") + if rule_parameters.get("check_encryption", "true").lower() == "true" and not check_encryption(bucket_name): + issues.append("encryption") + if rule_parameters.get("check_access_logging", "true").lower() == "true" and not check_access_logging(bucket_name): + issues.append("access logging") + if rule_parameters.get("check_object_locking", "true").lower() == "true" and not check_object_locking(bucket_name): + issues.append("object locking") + if rule_parameters.get("check_versioning", "true").lower() == "true" and not check_versioning(bucket_name): + issues.append("versioning") + + return issues + + +def get_bucket_name_from_data_source(data_source: Union[dict, GetDataSourceResponseTypeDef]) -> str | None: + """Extract bucket name from data source configuration. + + Args: + data_source (Union[dict, GetDataSourceResponseTypeDef]): Data source configuration + + Returns: + str | None: Bucket name if found, None otherwise + """ + try: + if (("dataSource" in data_source + and "dataSourceConfiguration" in data_source["dataSource"] + and "s3Configuration" in data_source["dataSource"]["dataSourceConfiguration"])): + s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] + bucket_arn = s3_config.get("bucketArn", "") + + if not bucket_arn: + return None + + bucket_name = bucket_arn.split(":")[-1] + return bucket_name.split("/")[0] if "/" in bucket_name else bucket_name + except Exception as e: + LOGGER.error(f"Error processing data source: {str(e)}") + return None + + +def check_knowledge_base(kb_id: str, rule_parameters: dict) -> list[str]: + """Check a knowledge base's data sources for S3 bucket compliance. + + Args: + kb_id (str): Knowledge base ID + rule_parameters (dict): Rule parameters containing check flags + + Returns: + list[str]: List of non-compliant bucket messages + """ + non_compliant_buckets = [] + data_sources_paginator = bedrock_agent_client.get_paginator("list_data_sources") + + for ds_page in data_sources_paginator.paginate(knowledgeBaseId=kb_id): + for ds in ds_page.get("dataSourceSummaries", []): + data_source = bedrock_agent_client.get_data_source( + knowledgeBaseId=kb_id, + dataSourceId=ds["dataSourceId"] + ) + + bucket_name = get_bucket_name_from_data_source(data_source) + if not bucket_name: + continue + + issues = check_bucket_configuration(bucket_name, rule_parameters) + if issues: + non_compliant_buckets.append(f"{bucket_name} (missing: {', '.join(issues)})") + + return non_compliant_buckets + + def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: """Evaluate if Bedrock Knowledge Base S3 bucket has required configurations. @@ -40,117 +215,12 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: tuple[str, str]: Compliance status and annotation """ try: - # Get all knowledge bases non_compliant_buckets = [] paginator = bedrock_agent_client.get_paginator("list_knowledge_bases") - + for page in paginator.paginate(): for kb in page["knowledgeBaseSummaries"]: - kb_id = kb["knowledgeBaseId"] - - # List data sources for this knowledge base - data_sources_paginator = bedrock_agent_client.get_paginator("list_data_sources") - - for ds_page in data_sources_paginator.paginate(knowledgeBaseId=kb_id): - for ds in ds_page.get("dataSourceSummaries", []): - data_source = bedrock_agent_client.get_data_source( - knowledgeBaseId=kb_id, - dataSourceId=ds["dataSourceId"] - ) - - # Check if this is an S3 data source and extract bucket name - try: - # Use a custom JSON encoder to handle datetime objects - class DateTimeEncoder(json.JSONEncoder): - def default(self, obj): - if isinstance(obj, datetime): - return obj.isoformat() - return super().default(obj) - - LOGGER.info(f"Data source structure: {json.dumps(data_source, cls=DateTimeEncoder)}") - - if "dataSource" in data_source and "dataSourceConfiguration" in data_source["dataSource"]: - if "s3Configuration" in data_source["dataSource"]["dataSourceConfiguration"]: - s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] - bucket_arn = s3_config.get("bucketArn", "") - - if not bucket_arn: - LOGGER.info(f"No bucket ARN found for data source {ds['dataSourceId']}") - continue - - # Extract bucket name from ARN - # ARN format: arn:aws:s3:::bucket-name - bucket_name = bucket_arn.split(":")[-1] - - # If bucket name contains a path, extract just the bucket name - if "/" in bucket_name: - bucket_name = bucket_name.split("/")[0] - else: - continue - else: - continue - except Exception as e: - LOGGER.error(f"Error processing data source: {str(e)}") - continue - - LOGGER.info(f"Checking S3 bucket: {bucket_name}") - - issues = [] - - # Check retention - if rule_parameters.get("check_retention", "true").lower() == "true": - try: - lifecycle = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name) - if not any(rule.get("Expiration") for rule in lifecycle.get("Rules", [])): - issues.append("retention") - except ClientError as e: - if e.response["Error"]["Code"] == "NoSuchLifecycleConfiguration": - issues.append("retention") - elif e.response["Error"]["Code"] != "NoSuchBucket": - LOGGER.error(f"Error checking retention for bucket {bucket_name}: {str(e)}") - - # Check encryption - if rule_parameters.get("check_encryption", "true").lower() == "true": - try: - encryption = s3_client.get_bucket_encryption(Bucket=bucket_name) - if not encryption.get("ServerSideEncryptionConfiguration"): - issues.append("encryption") - except ClientError as e: - if e.response["Error"]["Code"] != "NoSuchBucket": - issues.append("encryption") - - # Check server access logging - if rule_parameters.get("check_access_logging", "true").lower() == "true": - try: - logging_config = s3_client.get_bucket_logging(Bucket=bucket_name) - if not logging_config.get("LoggingEnabled"): - issues.append("access logging") - except ClientError as e: - if e.response["Error"]["Code"] != "NoSuchBucket": - issues.append("access logging") - - # Check object lock - if rule_parameters.get("check_object_locking", "true").lower() == "true": - try: - lock_config = s3_client.get_object_lock_configuration(Bucket=bucket_name) - if not lock_config.get("ObjectLockConfiguration"): - issues.append("object locking") - except ClientError as e: - if e.response["Error"]["Code"] != "NoSuchBucket": - issues.append("object locking") - - # Check versioning - if rule_parameters.get("check_versioning", "true").lower() == "true": - try: - versioning = s3_client.get_bucket_versioning(Bucket=bucket_name) - if versioning.get("Status") != "Enabled": - issues.append("versioning") - except ClientError as e: - if e.response["Error"]["Code"] != "NoSuchBucket": - issues.append("versioning") - - if issues: - non_compliant_buckets.append(f"{bucket_name} (missing: {', '.join(issues)})") + 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)}" @@ -160,7 +230,8 @@ def default(self, obj): LOGGER.error(f"Error evaluating Knowledge Base S3 bucket configurations: {str(e)}") return "ERROR", f"Error evaluating compliance: {str(e)}" -def lambda_handler(event: dict, context: Any) -> None: + +def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 """Lambda handler. Args: @@ -186,7 +257,6 @@ def lambda_handler(event: dict, context: Any) -> None: LOGGER.info(f"Compliance evaluation result: {compliance_type}") LOGGER.info(f"Annotation: {annotation}") - config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) + config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) # type: ignore LOGGER.info("Compliance evaluation complete.") - \ No newline at end of file diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py index 7d2c83f46..9c95fe55d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py @@ -29,7 +29,47 @@ secretsmanager_client = boto3.client("secretsmanager", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) -def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: + +def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: CFQ004 + """Check if a knowledge base's vector store is using KMS encrypted secrets. + + Args: + kb_id (str): Knowledge base ID + kb_name (str): Knowledge base name + + Raises: + ClientError: If there is an error accessing the knowledge base or secret. + + Returns: + tuple[bool, str]: (is_compliant, message) + """ + try: + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) + vector_store = kb_details.get("vectorStoreConfiguration") + + if not vector_store or not isinstance(vector_store, dict): + return False, f"{kb_name} (no vector store configuration)" + + secret_arn = vector_store.get("secretArn") + if not secret_arn: + return False, f"{kb_name} (no secret configured)" + + try: + secret_details = secretsmanager_client.describe_secret(SecretId=secret_arn) + if not secret_details.get("KmsKeyId"): + return False, f"{kb_name} (secret not using CMK)" + return True, "" + except ClientError as e: + if e.response["Error"]["Code"] == "AccessDeniedException": + return False, f"{kb_name} (secret access denied)" + raise + except ClientError as e: + if e.response["Error"]["Code"] == "AccessDeniedException": + return False, f"{kb_name} (access denied)" + raise + + +def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 """Evaluate if Bedrock Knowledge Base vector stores are using KMS encrypted secrets. Args: @@ -41,41 +81,14 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: 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) - - try: - # Get knowledge base details - kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) - vector_store = kb_details.get("vectorStoreConfiguration") - - if vector_store: - secret_arn = vector_store.get("secretArn") - if not secret_arn: - non_compliant_kbs.append(f"{kb_name} (no secret configured)") - continue - - try: - # Check if secret uses CMK - secret_details = secretsmanager_client.describe_secret(SecretId=secret_arn) - if not secret_details.get("KmsKeyId"): - non_compliant_kbs.append(f"{kb_name} (secret not using CMK)") - except ClientError as e: - LOGGER.error(f"Error checking secret {secret_arn}: {str(e)}") - if e.response["Error"]["Code"] == "AccessDeniedException": - non_compliant_kbs.append(f"{kb_name} (secret access denied)") - else: - raise - - except ClientError as e: - LOGGER.error(f"Error checking knowledge base {kb_name}: {str(e)}") - if e.response["Error"]["Code"] == "AccessDeniedException": - non_compliant_kbs.append(f"{kb_name} (access denied)") - else: - raise + is_compliant, message = check_knowledge_base(kb_id, kb_name) + if not is_compliant: + non_compliant_kbs.append(message) if non_compliant_kbs: return "NON_COMPLIANT", f"The following knowledge bases have vector store secret issues: {'; '.join(non_compliant_kbs)}" @@ -85,7 +98,8 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: LOGGER.error(f"Error evaluating Bedrock Knowledge Base vector store secrets: {str(e)}") return "ERROR", f"Error evaluating compliance: {str(e)}" -def lambda_handler(event: dict, context: Any) -> None: + +def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 """Lambda handler. Args: @@ -111,6 +125,6 @@ def lambda_handler(event: dict, context: Any) -> None: LOGGER.info(f"Compliance evaluation result: {compliance_type}") LOGGER.info(f"Annotation: {annotation}") - config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) + config_client.put_evaluations(Evaluations=[evaluation], ResultToken=event["resultToken"]) # type: ignore - LOGGER.info("Compliance evaluation complete.") \ No newline at end of file + LOGGER.info("Compliance evaluation complete.") From bfbebb8691c122fcdb2afe2c1c23d61e827e238b Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 08:23:35 -0600 Subject: [PATCH 14/36] fixing more flake8 and mypy errors --- .../rules/sra_bedrock_check_kb_ingestion_encryption/app.py | 2 +- .../sra_bedrock_check_kb_opensearch_encryption/app.py | 6 +++--- .../lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py | 2 +- .../solutions/genai/bedrock_org/lambda/src/app.py | 7 ++++--- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 065a434fe..5cccee7b4 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py @@ -29,7 +29,7 @@ config_client = boto3.client("config", region_name=AWS_REGION) -def check_data_sources(kb_id: str, kb_name: str) -> str | None: # noqa: CFQ004 +def check_data_sources(kb_id: str, kb_name: str) -> str | None: # type: ignore # noqa: CFQ004 """Check if a knowledge base's data sources are encrypted. Args: diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index 76018167b..65ca2b0ea 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -31,7 +31,7 @@ config_client = boto3.client("config", region_name=AWS_REGION) -def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: +def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: # type: ignore """Check OpenSearch Serverless collection encryption. Args: @@ -60,7 +60,7 @@ def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: return None -def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # noqa: CFQ004 +def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # type: ignore # noqa: CFQ004 """Check standard OpenSearch domain encryption. Args: @@ -83,7 +83,7 @@ def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # no return None -def check_knowledge_base(kb_id: str, kb_name: str) -> str | None: # noqa: CFQ004 +def check_knowledge_base(kb_id: str, kb_name: str) -> str | None: # type: ignore # noqa: CFQ004 """Check a knowledge base's OpenSearch configuration. Args: diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index e4679bbea..affc0cc54 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -148,7 +148,7 @@ def check_bucket_configuration(bucket_name: str, rule_parameters: dict) -> list[ return issues -def get_bucket_name_from_data_source(data_source: Union[dict, GetDataSourceResponseTypeDef]) -> str | None: +def get_bucket_name_from_data_source(data_source: Union[dict, GetDataSourceResponseTypeDef]) -> str | None: # type: ignore """Extract bucket name from data source configuration. Args: diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 7261aa9fd..1796e1437 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -202,15 +202,16 @@ def load_sra_cloudwatch_dashboard() -> dict: + r'\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\]\}$', "SRA-BEDROCK-CHECK-KB-LOGGING": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', - "SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' - + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', + "SRA-BEDROCK-CHECK-KB-INGESTION-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*' + + r'"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', "SRA-BEDROCK-CHECK-KB-S3-BUCKET": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*' + r'\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*\{(\s*"check_retention"\s*:\s*"(true|false)")?(\s*,\s*"check_encryption"\s*:\s*' + r'"(true|false)")?(\s*,\s*"check_access_logging"\s*:\s*"(true|false)")?(\s*,\s*"check_object_locking"\s*:\s*"(true|false)")?(\s*,\s*' + r'"check_versioning"\s*:\s*"(true|false)")?\s*\}\}$', "SRA-BEDROCK-CHECK-KB-VECTOR-STORE-SECRET": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*' + r'"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', - "SRA-BEDROCK-CHECK-KB-OPENSEARCH-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', + "SRA-BEDROCK-CHECK-KB-OPENSEARCH-ENCRYPTION": r'^\{"deploy"\s*:\s*"(true|false)",\s*"accounts"\s*:\s*\[((?:"[0-9]+"(?:\s*,\s*)?)*)\],\s*' + + r'"regions"\s*:\s*\[((?:"[a-z0-9-]+"(?:\s*,\s*)?)*)\],\s*"input_params"\s*:\s*(\{\})\}$', } # Instantiate sra class objects From 7ace2ffc0df8a08bbf8e7be9fafd67dec37cab9c Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 08:28:38 -0600 Subject: [PATCH 15/36] fixing formatter errors --- .../rules/sra_bedrock_check_kb_logging/app.py | 16 +++++++--------- .../app.py | 8 ++------ .../rules/sra_bedrock_check_kb_s3_bucket/app.py | 14 +++++++------- 3 files changed, 16 insertions(+), 22 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index af8fcf725..d4f4de0bb 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -44,9 +44,9 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ0 try: # List all knowledge bases kb_list = [] - paginator = bedrock_agent_client.get_paginator('list_knowledge_bases') + paginator = bedrock_agent_client.get_paginator("list_knowledge_bases") for page in paginator.paginate(): - kb_list.extend(page.get('knowledgeBaseSummaries', [])) + kb_list.extend(page.get("knowledgeBaseSummaries", [])) if not kb_list: return "COMPLIANT", "No knowledge bases found in the account" @@ -55,20 +55,18 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ0 # Check each knowledge base for logging configuration for kb in kb_list: - kb_id = kb['knowledgeBaseId'] + kb_id = kb["knowledgeBaseId"] try: - kb_details = bedrock_agent_client.get_knowledge_base( - knowledgeBaseId=kb_id - ) + kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) # Check if logging is enabled - logging_config = kb_details.get('loggingConfiguration', {}) - if not isinstance(logging_config, dict) or not logging_config.get('enabled', False): + logging_config = kb_details.get("loggingConfiguration", {}) + if not isinstance(logging_config, dict) or not logging_config.get("enabled", False): non_compliant_kbs.append(f"{kb_id} ({kb.get('name', 'unnamed')})") except ClientError as e: LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") - if e.response['Error']['Code'] == 'AccessDeniedException': + if e.response["Error"]["Code"] == "AccessDeniedException": non_compliant_kbs.append(f"{kb_id} (access denied)") else: raise diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index 65ca2b0ea..2aa8659b6 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -42,10 +42,7 @@ def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: str | None: Error message if non-compliant, None if compliant """ try: - collection = opensearch_serverless_client.get_security_policy( - name=collection_id, - type="encryption" - ) + collection = opensearch_serverless_client.get_security_policy(name=collection_id, type="encryption") security_policy = collection.get("securityPolicyDetail", {}) if security_policy.get("Type") == "encryption": security_policies = security_policy.get("SecurityPolicies", []) @@ -150,8 +147,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 if non_compliant_kbs: return "NON_COMPLIANT", ( - "The following knowledge bases have OpenSearch vector stores not encrypted with CMK: " - + f"{'; '.join(non_compliant_kbs)}" + "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" diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index affc0cc54..c2cc455fa 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -11,6 +11,7 @@ import logging import os from typing import Any, Union + import boto3 from botocore.exceptions import ClientError from mypy_boto3_bedrock_agent.type_defs import GetDataSourceResponseTypeDef @@ -158,9 +159,11 @@ def get_bucket_name_from_data_source(data_source: Union[dict, GetDataSourceRespo str | None: Bucket name if found, None otherwise """ try: - if (("dataSource" in data_source - and "dataSourceConfiguration" in data_source["dataSource"] - and "s3Configuration" in data_source["dataSource"]["dataSourceConfiguration"])): + if ( + "dataSource" in data_source + and "dataSourceConfiguration" in data_source["dataSource"] + and "s3Configuration" in data_source["dataSource"]["dataSourceConfiguration"] + ): s3_config = data_source["dataSource"]["dataSourceConfiguration"]["s3Configuration"] bucket_arn = s3_config.get("bucketArn", "") @@ -189,10 +192,7 @@ def check_knowledge_base(kb_id: str, rule_parameters: dict) -> list[str]: for ds_page in data_sources_paginator.paginate(knowledgeBaseId=kb_id): for ds in ds_page.get("dataSourceSummaries", []): - data_source = bedrock_agent_client.get_data_source( - knowledgeBaseId=kb_id, - dataSourceId=ds["dataSourceId"] - ) + data_source = bedrock_agent_client.get_data_source(knowledgeBaseId=kb_id, dataSourceId=ds["dataSourceId"]) bucket_name = get_bucket_name_from_data_source(data_source) if not bucket_name: From ecb32f7ea902ff016d2af92c32036d41eda7126b Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 08:42:13 -0600 Subject: [PATCH 16/36] update pyproject.toml --- pyproject.toml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f8e9d4d0b..fbeec1fa4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,13 @@ reportGeneralTypeIssues = "none" reportTypedDictNotRequiredAccess = "none" [tool.pylic] -safe_licenses = ["MIT License", "BSD License", "Apache Software License"] +safe_licenses = [ + "MIT", + "BSD-2-Clause", + "Apache-2.0", + "ISC", + "Python-2.0" +] [tool.vulture] ignore_decorators = ["@helper.*"] From 4292ec4d6b0643063aeda91b81117c450f3c6aab Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 08:47:53 -0600 Subject: [PATCH 17/36] update pyproject.toml --- pyproject.toml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fbeec1fa4..ff7f5162b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -127,8 +127,11 @@ safe_licenses = [ "MIT", "BSD-2-Clause", "Apache-2.0", - "ISC", - "Python-2.0" + "MIT License", + "BSD License", + "Apache Software License", + "Python Software Foundation License", + "ISC License (ISCL)" ] [tool.vulture] From 1b7bca5e86193069dea9ca3874ded1ea5a57e29f Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 11:58:34 -0600 Subject: [PATCH 18/36] fixing errors and permissions --- .../app.py | 51 +++++++++++++------ .../sra_config_lambda_iam_permissions.json | 6 +++ 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index 2aa8659b6..c7d76fe2d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -72,7 +72,8 @@ 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)}") @@ -80,7 +81,7 @@ def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # ty 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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 0e169c31b..6315654a8 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/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": "*" } ] } From f26f2118f6a0fa0d09151847d1e68d6aebd2808c Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Tue, 25 Mar 2025 23:00:37 -0600 Subject: [PATCH 19/36] refactor for mypy fixes; updates for finding kmsarn in collection for kb --- .../app.py | 44 ++++++++++++++----- .../sra_config_lambda_iam_permissions.json | 6 ++- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index c7d76fe2d..9c93257a0 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py @@ -31,7 +31,7 @@ config_client = boto3.client("config", region_name=AWS_REGION) -def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: # type: ignore +def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: # type: ignore # noqa: CFQ004 """Check OpenSearch Serverless collection encryption. Args: @@ -42,19 +42,41 @@ def check_opensearch_serverless(collection_id: str, kb_name: str) -> str | None: str | None: Error message if non-compliant, None if compliant """ try: - collection = opensearch_serverless_client.get_security_policy(name=collection_id, type="encryption") - security_policy = collection.get("securityPolicyDetail", {}) - if security_policy.get("Type") == "encryption": - security_policies = security_policy.get("SecurityPolicies", []) - if isinstance(security_policies, list) and security_policies: - encryption_policy = security_policies[0] - kms_key_arn = encryption_policy.get("KmsARN", "") - if not kms_key_arn or "aws/opensearchserverless" in kms_key_arn: - return f"{kb_name} (OpenSearch Serverless not using CMK)" + # Get collection details to get the collection name + collection_response = opensearch_serverless_client.batch_get_collection(ids=[collection_id]) + LOGGER.info(f"Collection details: {json.dumps(collection_response, default=str)}") + + 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)" + + 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)" + + # Get the specific policy details using the collection name + policy_details = opensearch_serverless_client.get_security_policy(name=collection_name, type="encryption") + LOGGER.info(f"Policy details for {collection_name}: {json.dumps(policy_details, default=str)}") + + policy_details_dict = json.loads(json.dumps(policy_details, default=str)) + policy_details_dict = policy_details_dict.get("securityPolicyDetail", {}).get("policy", {}) + 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)" + + 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 None + except ClientError as e: LOGGER.error(f"Error checking OpenSearch Serverless collection: {str(e)}") return f"{kb_name} (error checking OpenSearch Serverless)" - return None def check_opensearch_domain(domain_name: str, kb_name: str) -> str | None: # type: ignore # noqa: CFQ004 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 6315654a8..0405782b3 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -228,7 +228,11 @@ { "Sid": "AllowOpenSearchServerlessAccess", "Effect": "Allow", - "Action": ["aoss:GetSecurityPolicy"], + "Action": [ + "aoss:GetSecurityPolicy", + "aoss:ListSecurityPolicies", + "aoss:BatchGetCollection" + ], "Resource": "*" } ] From e0a413a6a3fe34fc4079056049edc167601aa896 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Wed, 26 Mar 2025 11:52:15 -0600 Subject: [PATCH 20/36] fixing for changes made because of mypy --- .../lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index c2cc455fa..285672ef8 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py @@ -10,11 +10,10 @@ import json import logging import os -from typing import Any, Union +from typing import Any, Dict import boto3 from botocore.exceptions import ClientError -from mypy_boto3_bedrock_agent.type_defs import GetDataSourceResponseTypeDef # Setup Default Logger LOGGER = logging.getLogger(__name__) @@ -149,11 +148,11 @@ def check_bucket_configuration(bucket_name: str, rule_parameters: dict) -> list[ return issues -def get_bucket_name_from_data_source(data_source: Union[dict, GetDataSourceResponseTypeDef]) -> str | None: # type: ignore +def get_bucket_name_from_data_source(data_source: Dict[str, Any]) -> str | None: # type: ignore """Extract bucket name from data source configuration. Args: - data_source (Union[dict, GetDataSourceResponseTypeDef]): Data source configuration + data_source (Dict[str, Any]): Data source configuration Returns: str | None: Bucket name if found, None otherwise @@ -194,7 +193,7 @@ def check_knowledge_base(kb_id: str, rule_parameters: dict) -> list[str]: for ds in ds_page.get("dataSourceSummaries", []): data_source = bedrock_agent_client.get_data_source(knowledgeBaseId=kb_id, dataSourceId=ds["dataSourceId"]) - bucket_name = get_bucket_name_from_data_source(data_source) + bucket_name = get_bucket_name_from_data_source(data_source) # type: ignore if not bucket_name: continue From 9b01f3f9432ccb34bd080cdc3a4d9e2b3bde5e82 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Wed, 26 Mar 2025 17:18:01 -0600 Subject: [PATCH 21/36] updating for mypy fixes; also resolved API call issues --- .../rules/sra_bedrock_check_kb_logging/app.py | 99 +++++++++++++++---- .../sra_config_lambda_iam_permissions.json | 11 +++ 2 files changed, 91 insertions(+), 19 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index d4f4de0bb..e9711005a 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -10,10 +10,9 @@ import json import logging import os -from typing import Any +from typing import Any, Optional, Tuple import boto3 -from botocore.exceptions import ClientError # Setup Default Logger LOGGER = logging.getLogger(__name__) @@ -27,6 +26,77 @@ # Initialize AWS clients bedrock_agent_client = boto3.client("bedrock-agent", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) +logs_client = boto3.client("logs", region_name=AWS_REGION) +sts_client = boto3.client("sts", region_name=AWS_REGION) + + +def check_kb_logging(kb_id: str) -> Tuple[bool, Optional[str]]: + """Check if knowledge base has CloudWatch logging enabled. + + Args: + kb_id (str): Knowledge base ID + + Returns: + Tuple[bool, Optional[str]]: (True if logging is enabled, destination type if found) + """ + try: + account_id = sts_client.get_caller_identity()["Account"] + kb_arn = f"arn:aws:bedrock:{AWS_REGION}:{account_id}:knowledge-base/{kb_id}" + LOGGER.info(f"Checking logging for KB ARN: {kb_arn}") + + # Get delivery sources + delivery_sources = logs_client.describe_delivery_sources() + LOGGER.info(f"Found {len(delivery_sources.get('deliverySources', []))} delivery sources") + + for source in delivery_sources.get("deliverySources", []): + LOGGER.info(f"Checking source: {source.get('name')}") + if kb_arn in source.get("resourceArns", []): + source_name = source.get("name") + LOGGER.info(f"Found matching source name: {source_name}") + if not source_name: + continue + + # Get deliveries to find the delivery ID + LOGGER.info("Calling describe_deliveries API") + deliveries = logs_client.describe_deliveries() + LOGGER.info(f"Found {len(deliveries.get('deliveries', []))} deliveries") + + for delivery in deliveries.get("deliveries", []): + LOGGER.info(f"Checking delivery: {delivery.get('id')} with source name: {delivery.get('deliverySourceName')}") + if delivery.get("deliverySourceName") == source_name: + delivery_id = delivery.get("id") + LOGGER.info(f"Found matching delivery ID: {delivery_id}") + if not delivery_id: + continue + + # Get delivery details to get the destination ARN + LOGGER.info(f"Calling get_delivery API with ID: {delivery_id}") + delivery_details = logs_client.get_delivery(id=delivery_id) + LOGGER.info(f"Delivery details: {delivery_details}") + + delivery_destination_arn = delivery_details.get("delivery", {}).get("deliveryDestinationArn") + LOGGER.info(f"Found delivery destination ARN: {delivery_destination_arn}") + if not delivery_destination_arn: + continue + + # Get delivery destinations to match the ARN + LOGGER.info("Calling describe_delivery_destinations API") + delivery_destinations = logs_client.describe_delivery_destinations() + LOGGER.info(f"Found {len(delivery_destinations.get('deliveryDestinations', []))} delivery destinations") + + for destination in delivery_destinations.get("deliveryDestinations", []): + LOGGER.info(f"Checking destination: {destination.get('name')} with ARN: {destination.get('arn')}") + if destination.get("arn") == delivery_destination_arn: + destination_type = destination.get("deliveryDestinationType") + LOGGER.info(f"Found matching destination with type: {destination_type}") + return True, destination_type + + LOGGER.info("No matching logging configuration found") + return False, None + + except Exception as e: + LOGGER.error(f"Error checking logging for knowledge base {kb_id}: {str(e)}") + return False, None def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ004, U100 @@ -35,9 +105,6 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ0 Args: rule_parameters (dict): Rule parameters from AWS Config rule. - Raises: - ClientError: If there is an error checking the knowledge base - Returns: tuple[str, str]: Compliance type and annotation message. """ @@ -52,28 +119,22 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: CFQ0 return "COMPLIANT", "No knowledge bases found in the account" non_compliant_kbs = [] + compliant_kbs = [] # Check each knowledge base for logging configuration for kb in kb_list: kb_id = kb["knowledgeBaseId"] - try: - kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) - - # Check if logging is enabled - logging_config = kb_details.get("loggingConfiguration", {}) - if not isinstance(logging_config, dict) or not logging_config.get("enabled", False): - non_compliant_kbs.append(f"{kb_id} ({kb.get('name', 'unnamed')})") + kb_name = kb.get("name", "unnamed") - except ClientError as e: - LOGGER.error(f"Error checking knowledge base {kb_id}: {str(e)}") - if e.response["Error"]["Code"] == "AccessDeniedException": - non_compliant_kbs.append(f"{kb_id} (access denied)") - else: - raise + 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") + else: + compliant_kbs.append(f"{kb_id} ({kb_name}) - logging configured 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", "All knowledge bases have logging enabled" + return "COMPLIANT", f"All knowledge bases have logging enabled: {', '.join(compliant_kbs)}" except Exception as e: LOGGER.error(f"Error evaluating Bedrock Knowledge Base logging configuration: {str(e)}") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 0405782b3..91e521150 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -135,6 +135,17 @@ "bedrock:GetKnowledgeBase" ], "Resource": "*" + }, + { + "Sid": "AllowDescribeDeliverySourcesAndDestinations", + "Effect": "Allow", + "Action": [ + "logs:DescribeDeliverySources", + "logs:DescribeDeliveryDestinations", + "logs:DescribeDeliveries", + "logs:GetDelivery" + ], + "Resource": "*" } ] }, From dba9d0e4659c2190e153131d4431ebe263585fe9 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Wed, 26 Mar 2025 17:18:36 -0600 Subject: [PATCH 22/36] flake8 issue --- .../lambda/rules/sra_bedrock_check_kb_logging/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index e9711005a..5bde70e0d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -30,7 +30,7 @@ sts_client = boto3.client("sts", region_name=AWS_REGION) -def check_kb_logging(kb_id: str) -> Tuple[bool, Optional[str]]: +def check_kb_logging(kb_id: str) -> Tuple[bool, Optional[str]]: # noqa: CCR001 """Check if knowledge base has CloudWatch logging enabled. Args: From 578d1da587590596d281a0d47b7bfda167e40316 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Mar 2025 11:56:57 -0600 Subject: [PATCH 23/36] fixes from flake8 and mypy refactoring; update api calls for failures --- .../app.py | 35 +++++++++++++++---- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 5cccee7b4..275789a12 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py @@ -29,8 +29,8 @@ 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 - """Check if a knowledge base's data sources are encrypted. +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 @@ -44,18 +44,41 @@ def check_data_sources(kb_id: str, kb_name: str) -> str | None: # type: ignore """ 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 data sources response)" + unencrypted_sources = [] for source in data_sources.get("dataSourceSummaries", []): + LOGGER.info(f"Source: {source}") if not isinstance(source, dict): continue - encryption_config = source.get("serverSideEncryptionConfiguration", {}) - if not isinstance(encryption_config, dict) or not encryption_config.get("kmsKeyArn"): - unencrypted_sources.append(source.get("name", source["dataSourceId"])) + + # 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'])} (access denied)") + continue if unencrypted_sources: - return f"{kb_name} (unencrypted sources: {', '.join(unencrypted_sources)})" + return f"{kb_name} (sources without KMS encryption: {', '.join(unencrypted_sources)})" return None except ClientError as e: LOGGER.error(f"Error checking data sources for knowledge base {kb_name}: {str(e)}") From 71ecfd6e2d0e814565a56bc008100a5dd77564c3 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Mar 2025 17:39:30 -0600 Subject: [PATCH 24/36] fixes for refactoring due to mypy and flake8; update API calls to appropriately get info --- .../app.py | 78 +++++++++++++++---- 1 file changed, 63 insertions(+), 15 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py index 9c95fe55d..5fccabacb 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py @@ -31,7 +31,7 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: CFQ004 - """Check if a knowledge base's vector store is using KMS encrypted secrets. + """Check if a knowledge base's vector store is using AWS Secrets Manager for credentials. Args: kb_id (str): Knowledge base ID @@ -45,27 +45,63 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: """ try: kb_details = bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) - vector_store = kb_details.get("vectorStoreConfiguration") - - if not vector_store or not isinstance(vector_store, dict): - return False, f"{kb_name} (no vector store configuration)" - - secret_arn = vector_store.get("secretArn") + LOGGER.info(f"KB Details: {json.dumps(kb_details, default=str)}") + + # Get the knowledge base object from the response + kb = kb_details.get("knowledgeBase", {}) + storage_config = kb.get("storageConfiguration") + LOGGER.info(f"Storage config from kb: {json.dumps(storage_config, default=str)}") + + if not storage_config or not isinstance(storage_config, dict): + return False, f"{kb_name} (Vector store configuration missing)" + + storage_type = storage_config.get("type") + LOGGER.info(f"Storage type: {storage_type}") + if not storage_type: + return False, f"{kb_name} (Vector store type not specified)" + + # Check if storage type is one of the supported types + supported_types = { + "PINECONE": "pineconeConfiguration", + "MONGO_DB_ATLAS": "mongoDbAtlasConfiguration", + "REDIS_ENTERPRISE_CLOUD": "redisEnterpriseCloudConfiguration", + "RDS": "rdsConfiguration" + } + + # If storage type is not supported, it's compliant (no credentials needed) + if storage_type not in supported_types: + LOGGER.info(f"Storage type {storage_type} not supported - no credentials needed") + return True, f"{kb_name} (Using unsupported vector store type '{storage_type}' - no credentials required)" + + # Get the configuration block for the storage type + config_key = supported_types[storage_type] + LOGGER.info(f"Config key: {config_key}") + type_config = storage_config.get(config_key) + LOGGER.info(f"Type config: {type_config}") + + if not type_config or not isinstance(type_config, dict): + return False, f"{kb_name} (Missing configuration for {storage_type} vector store)" + + # Check for credentials secret ARN + secret_arn = type_config.get("credentialsSecretArn") + LOGGER.info(f"Secret ARN: {secret_arn}") if not secret_arn: - return False, f"{kb_name} (no secret configured)" + return False, f"{kb_name} (Missing credentials secret for {storage_type} vector store)" try: + # Verify the secret exists and is using KMS encryption secret_details = secretsmanager_client.describe_secret(SecretId=secret_arn) + LOGGER.info(f"Secret details: {secret_details}") if not secret_details.get("KmsKeyId"): - return False, f"{kb_name} (secret not using CMK)" - return True, "" + return False, f"{kb_name} (Credentials secret for {storage_type} vector store not using CMK encryption)" + return True, f"{kb_name} (Using {storage_type} vector store with CMK-encrypted credentials)" except ClientError as e: if e.response["Error"]["Code"] == "AccessDeniedException": - return False, f"{kb_name} (secret access denied)" + return False, f"{kb_name} (Access denied to credentials secret for {storage_type} vector store)" raise except ClientError as e: if e.response["Error"]["Code"] == "AccessDeniedException": - return False, f"{kb_name} (access denied)" + return False, f"{kb_name} (Access denied to knowledge base)" raise @@ -80,19 +116,31 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 """ try: non_compliant_kbs = [] + 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"] + LOGGER.info(f"KB ID: {kb_id}") kb_name = kb.get("name", kb_id) + LOGGER.info(f"KB Name: {kb_name}") is_compliant, message = check_knowledge_base(kb_id, kb_name) - if not is_compliant: + if is_compliant: + compliant_kbs.append(message) + else: non_compliant_kbs.append(message) if non_compliant_kbs: - return "NON_COMPLIANT", f"The following knowledge bases have vector store secret issues: {'; '.join(non_compliant_kbs)}" - return "COMPLIANT", "All knowledge base vector stores are using KMS encrypted secrets" + return "NON_COMPLIANT", ( + "Knowledge base vector store compliance check results:\n" + + f"Compliant: {'; '.join(compliant_kbs)}\n" + + f"Non-compliant: {'; '.join(non_compliant_kbs)}" + ) + return "COMPLIANT", ( + "Knowledge base vector store compliance check results:\n" + + f"Compliant: {'; '.join(compliant_kbs)}" + ) except Exception as e: LOGGER.error(f"Error evaluating Bedrock Knowledge Base vector store secrets: {str(e)}") From 0e65503f8417a7cc5a327d0b0d5d122889ae148b Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Mar 2025 17:50:17 -0600 Subject: [PATCH 25/36] fix annotation language --- .../sra_bedrock_check_kb_ingestion_encryption/app.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 275789a12..576890879 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py @@ -78,7 +78,7 @@ def check_data_sources(kb_id: str, kb_name: str) -> str | None: # type: ignore continue if unencrypted_sources: - return f"{kb_name} (sources without KMS encryption: {', '.join(unencrypted_sources)})" + return f"{kb_name} (sources using default AWS-managed key instead of Customer Managed Key: {', '.join(unencrypted_sources)})" return None except ClientError as e: LOGGER.error(f"Error checking data sources for knowledge base {kb_name}: {str(e)}") @@ -109,8 +109,12 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 non_compliant_kbs.append(error) if non_compliant_kbs: - return "NON_COMPLIANT", f"The following knowledge bases have unencrypted data sources: {'; '.join(non_compliant_kbs)}" - return "COMPLIANT", "All knowledge base data sources are encrypted with KMS" + msg = ( + "The following knowledge bases are using default AWS-managed keys " + + f"instead of Customer Managed Keys: {'; '.join(non_compliant_kbs)}" + ) + return "NON_COMPLIANT", msg + return "COMPLIANT", "All knowledge base data sources are encrypted with Customer Managed Keys" except Exception as e: LOGGER.error(f"Error evaluating Bedrock Knowledge Base encryption: {str(e)}") From a81b013453216c293ac2572f69433974edc8ec2c Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Thu, 27 Mar 2025 17:57:31 -0600 Subject: [PATCH 26/36] black formatter changes --- .../app.py | 1 + .../sra_bedrock_check_eval_job_bucket/app.py | 1 + .../app.py | 1 + .../rules/sra_bedrock_check_guardrails/app.py | 1 + .../sra_bedrock_check_iam_user_access/app.py | 1 + .../app.py | 1 + .../app.py | 1 + .../app.py | 6 +- .../rules/sra_bedrock_check_kb_logging/app.py | 1 + .../app.py | 1 + .../sra_bedrock_check_kb_s3_bucket/app.py | 1 + .../app.py | 8 +- .../sra_bedrock_check_s3_endpoints/app.py | 1 + .../sra_bedrock_check_vpc_endpoints/app.py | 1 + .../genai/bedrock_org/lambda/src/app.py | 73 ++++++++++--------- .../bedrock_org/lambda/src/cfnresponse.py | 1 + .../genai/bedrock_org/lambda/src/sra_iam.py | 1 + .../genai/bedrock_org/lambda/src/sra_kms.py | 1 + .../bedrock_org/lambda/src/sra_lambda.py | 1 + .../genai/bedrock_org/lambda/src/sra_repo.py | 1 + .../genai/bedrock_org/lambda/src/sra_s3.py | 1 + .../genai/bedrock_org/lambda/src/sra_sns.py | 1 + .../bedrock_org/lambda/src/sra_ssm_params.py | 1 + .../genai/bedrock_org/lambda/src/sra_sts.py | 1 + 24 files changed, 63 insertions(+), 45 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_cloudwatch_endpoints/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_cloudwatch_endpoints/app.py index 6fe35ceb7..84ccd76bc 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_cloudwatch_endpoints/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_eval_job_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_eval_job_bucket/app.py index 5abfdf2c4..aed334a03 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_eval_job_bucket/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrail_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrail_encryption/app.py index c1113eb03..2b68f892d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrail_encryption/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrails/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrails/app.py index 72d0a7266..f9c37c3ce 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_guardrails/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_iam_user_access/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_iam_user_access/app.py index 10361ac68..1595681ed 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_iam_user_access/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_cloudwatch/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_cloudwatch/app.py index f94cc1fc4..570b5cfc4 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_cloudwatch/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_s3/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_s3/app.py index a88d73e12..608afef16 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_invocation_log_s3/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 576890879..61c71648f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_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 @@ -56,10 +57,7 @@ def check_data_sources(kb_id: str, kb_name: str) -> str | None: # type: ignore # Get the detailed data source configuration try: - source_details = bedrock_agent_client.get_data_source( - knowledgeBaseId=kb_id, - dataSourceId=source["dataSourceId"] - ) + 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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index 5bde70e0d..09619a96d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index 9c93257a0..2bef3bcbf 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index 285672ef8..d45c2704d 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py index 5fccabacb..dc551a910 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/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 @@ -65,7 +66,7 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: "PINECONE": "pineconeConfiguration", "MONGO_DB_ATLAS": "mongoDbAtlasConfiguration", "REDIS_ENTERPRISE_CLOUD": "redisEnterpriseCloudConfiguration", - "RDS": "rdsConfiguration" + "RDS": "rdsConfiguration", } # If storage type is not supported, it's compliant (no credentials needed) @@ -137,10 +138,7 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 + f"Compliant: {'; '.join(compliant_kbs)}\n" + f"Non-compliant: {'; '.join(non_compliant_kbs)}" ) - return "COMPLIANT", ( - "Knowledge base vector store compliance check results:\n" - + f"Compliant: {'; '.join(compliant_kbs)}" - ) + return "COMPLIANT", ("Knowledge base vector store compliance check results:\n" + f"Compliant: {'; '.join(compliant_kbs)}") except Exception as e: LOGGER.error(f"Error evaluating Bedrock Knowledge Base vector store secrets: {str(e)}") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_s3_endpoints/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_s3_endpoints/app.py index 7eb7df0f1..9d2965105 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_s3_endpoints/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_s3_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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_vpc_endpoints/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_vpc_endpoints/app.py index 234ba9ddb..aca1c0dd4 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_vpc_endpoints/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_vpc_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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 1796e1437..40c123dc9 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -8,6 +8,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + import copy import json import logging @@ -1117,9 +1118,9 @@ def deploy_metric_filters_and_alarms(region: str, accounts: list, resource_prope DRY_RUN_DATA[f"{filter_name}_CloudWatch_Alarm"] = "DRY_RUN: Deploy CloudWatch metric alarm" else: LOGGER.info(f"DRY_RUN: Filter deploy parameter is 'false'; Skip {filter_name} CloudWatch metric filter deployment") - DRY_RUN_DATA[ - f"{filter_name}_CloudWatch" - ] = "DRY_RUN: Filter deploy parameter is 'false'; Skip CloudWatch metric filter deployment" + DRY_RUN_DATA[f"{filter_name}_CloudWatch"] = ( + "DRY_RUN: Filter deploy parameter is 'false'; Skip CloudWatch metric filter deployment" + ) def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR001, CFQ001, C901 @@ -1215,9 +1216,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 if DRY_RUN is False: xacct_role = iam.create_role(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, cloudwatch.CROSS_ACCOUNT_TRUST_POLICY, SOLUTION_NAME) xacct_role_arn = xacct_role["Role"]["Arn"] - LIVE_RUN_DATA[ - f"OAMCrossAccountRoleCreate_{bedrock_account}" - ] = f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + LIVE_RUN_DATA[f"OAMCrossAccountRoleCreate_{bedrock_account}"] = ( + f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + ) CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] += 1 LOGGER.info(f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") @@ -1233,9 +1234,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 cloudwatch.CROSS_ACCOUNT_ROLE_NAME, ) else: - DRY_RUN_DATA[ - f"OAMCrossAccountRoleCreate_{bedrock_account}" - ] = f"DRY_RUN: Create {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + DRY_RUN_DATA[f"OAMCrossAccountRoleCreate_{bedrock_account}"] = ( + f"DRY_RUN: Create {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + ) else: LOGGER.info( f"CloudWatch observability access manager {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} cross-account role found in {bedrock_account}" @@ -1266,17 +1267,17 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 LOGGER.info(f"Attaching {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}...") if DRY_RUN is False: iam.attach_policy(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, policy_arn) - LIVE_RUN_DATA[ - f"OamXacctRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}" - ] = f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + LIVE_RUN_DATA[f"OamXacctRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}"] = ( + f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + ) CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["configuration_changes"] += 1 LOGGER.info(f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}") else: - DRY_RUN_DATA[ - f"OAMCrossAccountRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}" - ] = f"DRY_RUN: Attach {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + DRY_RUN_DATA[f"OAMCrossAccountRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}"] = ( + f"DRY_RUN: Attach {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" + ) # 5e) OAM link in bedrock account cloudwatch.CWOAM_CLIENT = sts.assume_role(bedrock_account, sts.CONFIGURATION_ROLE, "oam", bedrock_region) @@ -1285,9 +1286,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 if DRY_RUN is False: LOGGER.info("CloudWatch observability access manager link not found, creating...") oam_link_arn = cloudwatch.create_oam_link(oam_sink_arn) - LIVE_RUN_DATA[ - f"OAMLinkCreate_{bedrock_account}_{bedrock_region}" - ] = f"Created CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" + LIVE_RUN_DATA[f"OAMLinkCreate_{bedrock_account}_{bedrock_region}"] = ( + f"Created CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" + ) CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] += 1 @@ -1296,9 +1297,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 add_state_table_record("oam", "implemented", "oam link", "link", oam_link_arn, bedrock_account, bedrock_region, "oam_link") else: LOGGER.info("DRY_RUN: CloudWatch observability access manager link not found, creating...") - DRY_RUN_DATA[ - f"OAMLinkCreate_{bedrock_account}" - ] = f"DRY_RUN: Create CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" + DRY_RUN_DATA[f"OAMLinkCreate_{bedrock_account}"] = ( + f"DRY_RUN: Create CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" + ) # Set link arn to default value (for dry run) oam_link_arn = f"arn:aws:cloudwatch::{bedrock_account}:link/arn" else: @@ -1560,15 +1561,15 @@ def delete_custom_config_iam_role(rule_name: str, acct: str) -> None: # noqa: C if DRY_RUN is False: LOGGER.info(f"Detaching {policy['PolicyName']} IAM policy from account {acct} in {region}") iam.detach_policy(rule_name, policy["PolicyArn"]) - LIVE_RUN_DATA[ - f"{rule_name}_{acct}_{region}_PolicyDetach" - ] = f"Detached {policy['PolicyName']} IAM policy from account {acct} in {region}" + LIVE_RUN_DATA[f"{rule_name}_{acct}_{region}_PolicyDetach"] = ( + f"Detached {policy['PolicyName']} IAM policy from account {acct} in {region}" + ) CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 else: LOGGER.info(f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}") - DRY_RUN_DATA[ - f"{rule_name}_{acct}_{region}_Delete" - ] = f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}" + DRY_RUN_DATA[f"{rule_name}_{acct}_{region}_Delete"] = ( + f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}" + ) else: LOGGER.info(f"No IAM policies attached to {rule_name} for account {acct} in {region}") @@ -1586,9 +1587,9 @@ def delete_custom_config_iam_role(rule_name: str, acct: str) -> None: # noqa: C remove_state_table_record(policy_arn) else: LOGGER.info(f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}") - DRY_RUN_DATA[ - f"{rule_name}_{acct}_{region}_PolicyDelete" - ] = f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}" + DRY_RUN_DATA[f"{rule_name}_{acct}_{region}_PolicyDelete"] = ( + f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}" + ) else: LOGGER.info(f"{rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region} does not exist.") @@ -1806,18 +1807,18 @@ def delete_event(event: dict, context: Any) -> None: # noqa: CFQ001, CCR001, C9 for policy in cross_account_policies: LOGGER.info(f"Detaching {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role...") iam.detach_policy(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, policy["PolicyArn"]) - LIVE_RUN_DATA[ - "OAMCrossAccountRolePolicyDetach" - ] = f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + LIVE_RUN_DATA["OAMCrossAccountRolePolicyDetach"] = ( + f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + ) CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["configuration_changes"] += 1 LOGGER.info(f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") else: for policy in cross_account_policies: LOGGER.info(f"DRY_RUN: Detaching {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role...") - DRY_RUN_DATA[ - "OAMCrossAccountRolePolicyDetach" - ] = f"DRY_RUN: Detach {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + DRY_RUN_DATA["OAMCrossAccountRolePolicyDetach"] = ( + f"DRY_RUN: Detach {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" + ) else: LOGGER.info(f"No policies attached to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/cfnresponse.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/cfnresponse.py index 60d40e54f..9173c3c39 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/cfnresponse.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/cfnresponse.py @@ -1,4 +1,5 @@ """Amazon CFNResponse Module.""" + # mypy: ignore-errors # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: MIT-0 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py index 18198ca87..f116dab66 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + from __future__ import annotations import json diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py index 0322a2bce..0f62b58bf 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + from __future__ import annotations import logging diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_lambda.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_lambda.py index eba755e0f..3f04e5229 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_lambda.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_lambda.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + from __future__ import annotations import logging diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_repo.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_repo.py index 3e308f382..36807b12f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_repo.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_repo.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + import logging import os import shutil diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_s3.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_s3.py index 1bb11f385..6a419999e 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_s3.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_s3.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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sns.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sns.py index 1cb2f99ac..56bf4250f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sns.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sns.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + from __future__ import annotations import json diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py index 5e87ad8f2..efd7abf10 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + from __future__ import annotations import logging diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py index f36959023..1db7e439b 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py @@ -7,6 +7,7 @@ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 """ + import logging import os from typing import Any From f7f6d84766c0852a009a4641f417818e42154bf0 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Fri, 28 Mar 2025 10:10:51 -0600 Subject: [PATCH 27/36] update perms --- .../lambda/src/sra_config_lambda_iam_permissions.json | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 91e521150..26ea79031 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -231,8 +231,7 @@ "Sid": "AllowOpenSearchAccess", "Effect": "Allow", "Action": [ - "opensearch:DescribeDomain", - "opensearch:GetSecurityPolicy" + "es:DescribeDomain" ], "Resource": "*" }, From 58d21d5ed531bd9a4f791da0cb12ffac499d8232 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Fri, 28 Mar 2025 11:33:28 -0600 Subject: [PATCH 28/36] ensuring annotations do not exceed 256 char limit --- .../app.py | 21 ++--- .../rules/sra_bedrock_check_kb_logging/app.py | 49 +++++++++-- .../app.py | 50 +++++++----- .../sra_bedrock_check_kb_s3_bucket/app.py | 29 ++++++- .../app.py | 81 ++++++++++++++----- .../sra_config_lambda_iam_permissions.json | 1 + 6 files changed, 173 insertions(+), 58 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py index 61c71648f..72b5e6d35 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_ingestion_encryption/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index 09619a96d..9be476170 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py index 2bef3bcbf..e97f7df73 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_opensearch_encryption/app.py +++ b/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", diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py index d45c2704d..17f2a79cf 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_s3_bucket/app.py +++ b/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 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py index dc551a910..6c8ef4407 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_vector_store_secret/app.py @@ -11,7 +11,7 @@ import json import logging import os -from typing import Any +from typing import Any, List, Optional import boto3 from botocore.exceptions import ClientError @@ -30,6 +30,9 @@ secretsmanager_client = boto3.client("secretsmanager", region_name=AWS_REGION) config_client = boto3.client("config", region_name=AWS_REGION) +# Maximum annotation length for AWS Config PutEvaluations API +MAX_ANNOTATION_LENGTH = 256 + def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: CFQ004 """Check if a knowledge base's vector store is using AWS Secrets Manager for credentials. @@ -54,12 +57,12 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: LOGGER.info(f"Storage config from kb: {json.dumps(storage_config, default=str)}") if not storage_config or not isinstance(storage_config, dict): - return False, f"{kb_name} (Vector store configuration missing)" + return False, f"{kb_name} (No vector config)" storage_type = storage_config.get("type") LOGGER.info(f"Storage type: {storage_type}") if not storage_type: - return False, f"{kb_name} (Vector store type not specified)" + return False, f"{kb_name} (No store type)" # Check if storage type is one of the supported types supported_types = { @@ -72,7 +75,7 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: # If storage type is not supported, it's compliant (no credentials needed) if storage_type not in supported_types: LOGGER.info(f"Storage type {storage_type} not supported - no credentials needed") - return True, f"{kb_name} (Using unsupported vector store type '{storage_type}' - no credentials required)" + return True, f"{kb_name} ({storage_type} - no creds needed)" # Get the configuration block for the storage type config_key = supported_types[storage_type] @@ -81,31 +84,73 @@ def check_knowledge_base(kb_id: str, kb_name: str) -> tuple[bool, str]: # noqa: LOGGER.info(f"Type config: {type_config}") if not type_config or not isinstance(type_config, dict): - return False, f"{kb_name} (Missing configuration for {storage_type} vector store)" + return False, f"{kb_name} (Missing {storage_type} config)" # Check for credentials secret ARN secret_arn = type_config.get("credentialsSecretArn") LOGGER.info(f"Secret ARN: {secret_arn}") if not secret_arn: - return False, f"{kb_name} (Missing credentials secret for {storage_type} vector store)" + return False, f"{kb_name} (Missing secret)" try: # Verify the secret exists and is using KMS encryption secret_details = secretsmanager_client.describe_secret(SecretId=secret_arn) LOGGER.info(f"Secret details: {secret_details}") if not secret_details.get("KmsKeyId"): - return False, f"{kb_name} (Credentials secret for {storage_type} vector store not using CMK encryption)" - return True, f"{kb_name} (Using {storage_type} vector store with CMK-encrypted credentials)" + return False, f"{kb_name} (Secret not using CMK)" + return True, f"{kb_name} (Uses CMK)" except ClientError as e: if e.response["Error"]["Code"] == "AccessDeniedException": - return False, f"{kb_name} (Access denied to credentials secret for {storage_type} vector store)" + return False, f"{kb_name} (Secret access denied)" raise except ClientError as e: if e.response["Error"]["Code"] == "AccessDeniedException": - return False, f"{kb_name} (Access denied to knowledge base)" + return False, f"{kb_name} (KB access denied)" raise +def format_annotation(compliance_type: str, compliant_kbs: List[str], non_compliant_kbs: Optional[List[str]] = None) -> str: + """Format annotation message and ensure it doesn't exceed 256 characters. + + Args: + compliance_type (str): Compliance status + compliant_kbs (List[str]): List of compliant knowledge bases + non_compliant_kbs (Optional[List[str]]): List of non-compliant knowledge bases + + Returns: + str: Formatted annotation message + """ + if compliance_type == "ERROR": + return compliant_kbs[0] # In this case, compliant_kbs contains the error message + + non_compliant_count = len(non_compliant_kbs) if non_compliant_kbs else 0 + message = "" + + # Start with a brief message + if compliance_type == "NON_COMPLIANT": + base_message = "KB vector store check: " + if non_compliant_kbs: + combined = "; ".join(non_compliant_kbs) + # If message would be too long, provide a count instead of details + if len(base_message + combined) > MAX_ANNOTATION_LENGTH: + message = f"{base_message}{non_compliant_count} non-compliant KBs. See logs for details." + else: + message = base_message + combined + else: # COMPLIANT + if len(compliant_kbs) > 3: + message = f"All {len(compliant_kbs)} KBs comply with vector store requirements." + elif len("; ".join(compliant_kbs)) > MAX_ANNOTATION_LENGTH: + message = f"{len(compliant_kbs)} KBs comply with vector store requirements." + else: + message = "; ".join(compliant_kbs) + + # Final check to ensure we don't exceed limit + if len(message) > MAX_ANNOTATION_LENGTH: + return f"See CloudWatch logs for details. Found {non_compliant_count} issues." + + return message + + def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 """Evaluate if Bedrock Knowledge Base vector stores are using KMS encrypted secrets. @@ -132,17 +177,17 @@ def evaluate_compliance(rule_parameters: dict) -> tuple[str, str]: # noqa: U100 else: non_compliant_kbs.append(message) + # Log full details for CloudWatch reference + LOGGER.info(f"Compliant KBs: {'; '.join(compliant_kbs)}") if non_compliant_kbs: - return "NON_COMPLIANT", ( - "Knowledge base vector store compliance check results:\n" - + f"Compliant: {'; '.join(compliant_kbs)}\n" - + f"Non-compliant: {'; '.join(non_compliant_kbs)}" - ) - return "COMPLIANT", ("Knowledge base vector store compliance check results:\n" + f"Compliant: {'; '.join(compliant_kbs)}") + LOGGER.info(f"Non-compliant KBs: {'; '.join(non_compliant_kbs)}") + return "NON_COMPLIANT", format_annotation("NON_COMPLIANT", compliant_kbs, non_compliant_kbs) + return "COMPLIANT", format_annotation("COMPLIANT", compliant_kbs) except Exception as e: - LOGGER.error(f"Error evaluating Bedrock Knowledge Base vector store secrets: {str(e)}") - return "ERROR", f"Error evaluating compliance: {str(e)}" + error_msg = f"Error evaluating compliance: {str(e)}" + LOGGER.error(error_msg) + return "ERROR", format_annotation("ERROR", [error_msg]) def lambda_handler(event: dict, context: Any) -> None: # noqa: U100 diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json index 26ea79031..f055d5b6f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_config_lambda_iam_permissions.json @@ -241,6 +241,7 @@ "Action": [ "aoss:GetSecurityPolicy", "aoss:ListSecurityPolicies", + "aoss:BatchGetCollection", "aoss:BatchGetCollection" ], "Resource": "*" From a5c627e8d1cafebbaaceeb0600799451c5821bbc Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Fri, 28 Mar 2025 12:16:48 -0600 Subject: [PATCH 29/36] black formatter changes --- .../rules/sra_bedrock_check_kb_logging/app.py | 2 +- .../genai/bedrock_org/lambda/src/app.py | 72 +++++++++---------- 2 files changed, 37 insertions(+), 37 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py index 9be476170..bbc95016f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/rules/sra_bedrock_check_kb_logging/app.py @@ -53,7 +53,7 @@ def truncate_annotation(message: str) -> str: available_chars = MAX_ANNOTATION_LENGTH - len(reference) # Truncate message and add reference - truncated = message[:available_chars - 3] + "..." + truncated = message[: available_chars - 3] + "..." return truncated + reference diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 40c123dc9..8f2ac9322 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -1118,9 +1118,9 @@ def deploy_metric_filters_and_alarms(region: str, accounts: list, resource_prope DRY_RUN_DATA[f"{filter_name}_CloudWatch_Alarm"] = "DRY_RUN: Deploy CloudWatch metric alarm" else: LOGGER.info(f"DRY_RUN: Filter deploy parameter is 'false'; Skip {filter_name} CloudWatch metric filter deployment") - DRY_RUN_DATA[f"{filter_name}_CloudWatch"] = ( - "DRY_RUN: Filter deploy parameter is 'false'; Skip CloudWatch metric filter deployment" - ) + DRY_RUN_DATA[ + f"{filter_name}_CloudWatch" + ] = "DRY_RUN: Filter deploy parameter is 'false'; Skip CloudWatch metric filter deployment" def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR001, CFQ001, C901 @@ -1216,9 +1216,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 if DRY_RUN is False: xacct_role = iam.create_role(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, cloudwatch.CROSS_ACCOUNT_TRUST_POLICY, SOLUTION_NAME) xacct_role_arn = xacct_role["Role"]["Arn"] - LIVE_RUN_DATA[f"OAMCrossAccountRoleCreate_{bedrock_account}"] = ( - f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" - ) + LIVE_RUN_DATA[ + f"OAMCrossAccountRoleCreate_{bedrock_account}" + ] = f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] += 1 LOGGER.info(f"Created {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") @@ -1234,9 +1234,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 cloudwatch.CROSS_ACCOUNT_ROLE_NAME, ) else: - DRY_RUN_DATA[f"OAMCrossAccountRoleCreate_{bedrock_account}"] = ( - f"DRY_RUN: Create {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" - ) + DRY_RUN_DATA[ + f"OAMCrossAccountRoleCreate_{bedrock_account}" + ] = f"DRY_RUN: Create {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" else: LOGGER.info( f"CloudWatch observability access manager {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} cross-account role found in {bedrock_account}" @@ -1267,17 +1267,17 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 LOGGER.info(f"Attaching {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}...") if DRY_RUN is False: iam.attach_policy(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, policy_arn) - LIVE_RUN_DATA[f"OamXacctRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}"] = ( - f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" - ) + LIVE_RUN_DATA[ + f"OamXacctRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}" + ] = f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["configuration_changes"] += 1 LOGGER.info(f"Attached {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}") else: - DRY_RUN_DATA[f"OAMCrossAccountRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}"] = ( - f"DRY_RUN: Attach {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" - ) + DRY_RUN_DATA[ + f"OAMCrossAccountRolePolicyAttach_{policy_arn.split('/')[1]}_{bedrock_account}" + ] = f"DRY_RUN: Attach {policy_arn} policy to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role in {bedrock_account}" # 5e) OAM link in bedrock account cloudwatch.CWOAM_CLIENT = sts.assume_role(bedrock_account, sts.CONFIGURATION_ROLE, "oam", bedrock_region) @@ -1286,9 +1286,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 if DRY_RUN is False: LOGGER.info("CloudWatch observability access manager link not found, creating...") oam_link_arn = cloudwatch.create_oam_link(oam_sink_arn) - LIVE_RUN_DATA[f"OAMLinkCreate_{bedrock_account}_{bedrock_region}"] = ( - f"Created CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" - ) + LIVE_RUN_DATA[ + f"OAMLinkCreate_{bedrock_account}_{bedrock_region}" + ] = f"Created CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] += 1 @@ -1297,9 +1297,9 @@ def deploy_central_cloudwatch_observability(event: dict) -> None: # noqa: CCR00 add_state_table_record("oam", "implemented", "oam link", "link", oam_link_arn, bedrock_account, bedrock_region, "oam_link") else: LOGGER.info("DRY_RUN: CloudWatch observability access manager link not found, creating...") - DRY_RUN_DATA[f"OAMLinkCreate_{bedrock_account}"] = ( - f"DRY_RUN: Create CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" - ) + DRY_RUN_DATA[ + f"OAMLinkCreate_{bedrock_account}" + ] = f"DRY_RUN: Create CloudWatch observability access manager link in {bedrock_account} in {bedrock_region}" # Set link arn to default value (for dry run) oam_link_arn = f"arn:aws:cloudwatch::{bedrock_account}:link/arn" else: @@ -1561,15 +1561,15 @@ def delete_custom_config_iam_role(rule_name: str, acct: str) -> None: # noqa: C if DRY_RUN is False: LOGGER.info(f"Detaching {policy['PolicyName']} IAM policy from account {acct} in {region}") iam.detach_policy(rule_name, policy["PolicyArn"]) - LIVE_RUN_DATA[f"{rule_name}_{acct}_{region}_PolicyDetach"] = ( - f"Detached {policy['PolicyName']} IAM policy from account {acct} in {region}" - ) + LIVE_RUN_DATA[ + f"{rule_name}_{acct}_{region}_PolicyDetach" + ] = f"Detached {policy['PolicyName']} IAM policy from account {acct} in {region}" CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 else: LOGGER.info(f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}") - DRY_RUN_DATA[f"{rule_name}_{acct}_{region}_Delete"] = ( - f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}" - ) + DRY_RUN_DATA[ + f"{rule_name}_{acct}_{region}_Delete" + ] = f"DRY_RUN: Detach {policy['PolicyName']} IAM policy from account {acct} in {region}" else: LOGGER.info(f"No IAM policies attached to {rule_name} for account {acct} in {region}") @@ -1587,9 +1587,9 @@ def delete_custom_config_iam_role(rule_name: str, acct: str) -> None: # noqa: C remove_state_table_record(policy_arn) else: LOGGER.info(f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}") - DRY_RUN_DATA[f"{rule_name}_{acct}_{region}_PolicyDelete"] = ( - f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}" - ) + DRY_RUN_DATA[ + f"{rule_name}_{acct}_{region}_PolicyDelete" + ] = f"DRY_RUN: Delete {rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region}" else: LOGGER.info(f"{rule_name}-lamdba-basic-execution IAM policy for account {acct} in {region} does not exist.") @@ -1807,18 +1807,18 @@ def delete_event(event: dict, context: Any) -> None: # noqa: CFQ001, CCR001, C9 for policy in cross_account_policies: LOGGER.info(f"Detaching {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role...") iam.detach_policy(cloudwatch.CROSS_ACCOUNT_ROLE_NAME, policy["PolicyArn"]) - LIVE_RUN_DATA["OAMCrossAccountRolePolicyDetach"] = ( - f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" - ) + LIVE_RUN_DATA[ + "OAMCrossAccountRolePolicyDetach" + ] = f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1 CFN_RESPONSE_DATA["deployment_info"]["configuration_changes"] += 1 LOGGER.info(f"Detached {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") else: for policy in cross_account_policies: LOGGER.info(f"DRY_RUN: Detaching {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role...") - DRY_RUN_DATA["OAMCrossAccountRolePolicyDetach"] = ( - f"DRY_RUN: Detach {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" - ) + DRY_RUN_DATA[ + "OAMCrossAccountRolePolicyDetach" + ] = f"DRY_RUN: Detach {policy['PolicyArn']} policy from {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role" else: LOGGER.info(f"No policies attached to {cloudwatch.CROSS_ACCOUNT_ROLE_NAME} IAM role") From 3f1d91da3d99af1c4a6c5fd6dbd04ebd87216a7d Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Mon, 7 Apr 2025 10:39:08 -0600 Subject: [PATCH 30/36] ensuring rule delete operations don't target accounts not in rule_accounts (e.g. mgmt acct) --- aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 8f2ac9322..1145e4ee2 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -819,6 +819,9 @@ def deploy_config_rules(region: str, accounts: list, resource_properties: dict) if rule_deploy is False: LOGGER.info(f"{rule_name} is not to be deployed. Checking to see if it needs to be removed...") + if acct not in rule_accounts: + LOGGER.info(f"{rule_name} does not apply to {acct}; skipping attempt to delete...") + continue delete_custom_config_rule(rule_name, acct, region) delete_custom_config_iam_role(rule_name, acct) continue From db4a965ccd63f5848eb94b85bc7123bbe497ef87 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Mon, 7 Apr 2025 11:52:04 -0600 Subject: [PATCH 31/36] adding error handling with IAM policy deletes (race condition found) --- .../genai/bedrock_org/lambda/src/sra_iam.py | 33 ++++++++++++++++--- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py index f116dab66..891390231 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py @@ -165,7 +165,15 @@ def detach_policy(self, role_name: str, policy_arn: str) -> EmptyResponseMetadat Empty response metadata """ self.LOGGER.info("Detaching policy from %s.", role_name) - return self.IAM_CLIENT.detach_role_policy(RoleName=role_name, PolicyArn=policy_arn) + try: + response = self.IAM_CLIENT.detach_role_policy(RoleName=role_name, PolicyArn=policy_arn) + except ClientError as error: + if error.response["Error"]["Code"] == "NoSuchEntity": + self.LOGGER.info(f"Policy '{policy_arn}' is not attached to role '{role_name}'.") + else: + self.LOGGER.error(f"Error detaching policy '{policy_arn}' from role '{role_name}': {error}") + raise ValueError(f"Error detaching policy '{policy_arn}' from role '{role_name}': {error}") from None + return response def delete_policy(self, policy_arn: str) -> EmptyResponseMetadataTypeDef: """Delete IAM Policy. @@ -184,10 +192,25 @@ def delete_policy(self, policy_arn: str) -> EmptyResponseMetadataTypeDef: for version in page["Versions"]: if not version["IsDefaultVersion"]: self.LOGGER.info(f"Deleting policy version {version['VersionId']}") - self.IAM_CLIENT.delete_policy_version(PolicyArn=policy_arn, VersionId=version["VersionId"]) - sleep(1) - self.LOGGER.info("Policy version deleted.") - return self.IAM_CLIENT.delete_policy(PolicyArn=policy_arn) + try: + self.IAM_CLIENT.delete_policy_version(PolicyArn=policy_arn, VersionId=version["VersionId"]) + sleep(1) + self.LOGGER.info("Policy version deleted.") + except ClientError as error: + if error.response["Error"]["Code"] == "NoSuchEntity": + self.LOGGER.info(f"Policy version {version['VersionId']} not found.") + else: + self.LOGGER.error(f"Error deleting policy version {version['VersionId']}: {error}") + raise ValueError(f"Error deleting policy version {version['VersionId']}: {error}") from None + try: + response = self.IAM_CLIENT.delete_policy(PolicyArn=policy_arn) + except ClientError as error: + if error.response["Error"]["Code"] == "NoSuchEntity": + self.LOGGER.info(f"Policy {policy_arn} not found.") + else: + self.LOGGER.error(f"Error deleting policy {policy_arn}: {error}") + raise ValueError(f"Error deleting policy {policy_arn}: {error}") from None + return response def delete_role(self, role_name: str) -> EmptyResponseMetadataTypeDef: """Delete IAM role. From 36bc0b48f661c086d3ffde2ba2faead0f982f7c8 Mon Sep 17 00:00:00 2001 From: liamschn <57731583+liamschn@users.noreply.github.com> Date: Mon, 7 Apr 2025 11:53:52 -0600 Subject: [PATCH 32/36] handling flake8 issues --- .../solutions/genai/bedrock_org/lambda/src/sra_iam.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py index 891390231..97f9a3853 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_iam.py @@ -161,6 +161,9 @@ def detach_policy(self, role_name: str, policy_arn: str) -> EmptyResponseMetadat role_name: Name of the role for which the policy is removed from policy_arn: The Amazon Resource Name (ARN) of the policy to be detached + Raises: + ValueError: If an unexpected error occurs during the operation. + Returns: Empty response metadata """ @@ -175,12 +178,15 @@ def detach_policy(self, role_name: str, policy_arn: str) -> EmptyResponseMetadat raise ValueError(f"Error detaching policy '{policy_arn}' from role '{role_name}': {error}") from None return response - def delete_policy(self, policy_arn: str) -> EmptyResponseMetadataTypeDef: + def delete_policy(self, policy_arn: str) -> EmptyResponseMetadataTypeDef: # noqa: CCR001 """Delete IAM Policy. Args: policy_arn: The Amazon Resource Name (ARN) of the policy to be deleted + Raises: + ValueError: If an unexpected error occurs during the operation. + Returns: Empty response metadata """ From 7b7498d03c579e3c3d763bc2282974f0e9158085 Mon Sep 17 00:00:00 2001 From: cyphronix <57731583+liamschn@users.noreply.github.com> Date: Mon, 7 Apr 2025 15:22:08 -0600 Subject: [PATCH 33/36] fixing cleanup issues --- .../solutions/genai/bedrock_org/lambda/src/app.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index 1145e4ee2..adbe0578a 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -1880,8 +1880,8 @@ def delete_event(event: dict, context: Any) -> None: # noqa: CFQ001, CCR001, C9 for region in regions: delete_custom_config_rule(rule_name, acct, region) - # 5, 6, & 7) Detach IAM policies, delete IAM policy, delete IAM execution role for custom config rule lambda - delete_custom_config_iam_role(rule_name, acct) + # 5, 6, & 7) Detach IAM policies, delete IAM policy, delete IAM execution role for custom config rule lambda + delete_custom_config_iam_role(rule_name, acct) # Must infer the execution role arn because the function is being reported as non-existent at this point execution_role_arn = f"arn:aws:iam::{sts.MANAGEMENT_ACCOUNT}:role/{SOLUTION_NAME}-lambda" LOGGER.info(f"Removing state table record for lambda IAM execution role: {execution_role_arn}") From 6bd16bfdd86965bdee480467248e3180ed2666ad Mon Sep 17 00:00:00 2001 From: cyphronix <57731583+liamschn@users.noreply.github.com> Date: Wed, 9 Apr 2025 16:39:01 -0600 Subject: [PATCH 34/36] don't record CW centrol observability resources in dry-run if deploy is false --- .../genai/bedrock_org/lambda/src/app.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index adbe0578a..bedd61b70 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -1455,13 +1455,17 @@ def create_event(event: dict, context: Any) -> str: create_sns_messages(accounts, regions, topic_arn, event["ResourceProperties"], "configure") LOGGER.info(f"CFN_RESPONSE_DATA POST create_sns_messages: {CFN_RESPONSE_DATA}") - # 5) Central CloudWatch Observability (regional) - deploy_central_cloudwatch_observability(event) - LOGGER.info(f"CFN_RESPONSE_DATA POST deploy_central_cloudwatch_observability: {CFN_RESPONSE_DATA}") - - # 6) Cloudwatch dashboard in security account (home region, security account) - deploy_cloudwatch_dashboard(event) - LOGGER.info(f"CFN_RESPONSE_DATA POST deploy_cloudwatch_dashboard: {CFN_RESPONSE_DATA}") + central_observability_params = json.loads(event["ResourceProperties"]["SRA-BEDROCK-CENTRAL-OBSERVABILITY"]) + if central_observability_params["deploy"] is True: + # 5) Central CloudWatch Observability (regional) + deploy_central_cloudwatch_observability(event) + LOGGER.info(f"CFN_RESPONSE_DATA POST deploy_central_cloudwatch_observability: {CFN_RESPONSE_DATA}") + + # 6) Cloudwatch dashboard in security account (home region, security account) + deploy_cloudwatch_dashboard(event) + LOGGER.info(f"CFN_RESPONSE_DATA POST deploy_cloudwatch_dashboard: {CFN_RESPONSE_DATA}") + else: + LOGGER.info("CloudWatch observability deploy set to false, skipping deployment...") # End if DRY_RUN is False: From 99d1bd52e71a810ec56b228de5b609551e82dedf Mon Sep 17 00:00:00 2001 From: cyphronix <57731583+liamschn@users.noreply.github.com> Date: Wed, 9 Apr 2025 17:11:00 -0600 Subject: [PATCH 35/36] update condition statement --- aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py index bedd61b70..f0f0e893f 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py @@ -1456,7 +1456,7 @@ def create_event(event: dict, context: Any) -> str: LOGGER.info(f"CFN_RESPONSE_DATA POST create_sns_messages: {CFN_RESPONSE_DATA}") central_observability_params = json.loads(event["ResourceProperties"]["SRA-BEDROCK-CENTRAL-OBSERVABILITY"]) - if central_observability_params["deploy"] is True: + if central_observability_params["deploy"] == "true": # 5) Central CloudWatch Observability (regional) deploy_central_cloudwatch_observability(event) LOGGER.info(f"CFN_RESPONSE_DATA POST deploy_central_cloudwatch_observability: {CFN_RESPONSE_DATA}") From 98fffa08a5d4648f71d4857286b08a0fd8339049 Mon Sep 17 00:00:00 2001 From: cyphronix <57731583+liamschn@users.noreply.github.com> Date: Thu, 10 Apr 2025 12:32:06 -0600 Subject: [PATCH 36/36] enable key rotation for CMKs --- .../solutions/genai/bedrock_org/lambda/src/sra_kms.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py index 0f62b58bf..d86cfdaec 100644 --- a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py +++ b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_kms.py @@ -71,7 +71,13 @@ def create_kms_key(self, kms_client: KMSClient, key_policy: str, description: st KeyUsage="ENCRYPT_DECRYPT", CustomerMasterKeySpec="SYMMETRIC_DEFAULT", ) - return key_response["KeyMetadata"]["KeyId"] + key_id = key_response["KeyMetadata"]["KeyId"] + + # Enable key rotation + self.LOGGER.info(f"Enabling key rotation for key: {key_id}") + kms_client.enable_key_rotation(KeyId=key_id) + + return key_id def create_alias(self, kms_client: KMSClient, alias_name: str, target_key_id: str) -> None: """Create KMS alias.