adding sqs queue record to state table

ievgeniia ieromenko · ievgeniia ieromenko · commit dea7687ea98e · 2025-01-27T13:24:39.000-05:00
diff --git a/aws_sra_examples/solutions/genai/bedrock_guardrails/README.md b/aws_sra_examples/solutions/genai/bedrock_guardrails/README.md
@@ -1,4 +1,4 @@
-# SRA Bedrock Organizations Solution
+# SRA Bedrock Guardrails Solution
 
 ## Table of Contents
 - [Introduction](#introduction)
diff --git a/aws_sra_examples/solutions/genai/bedrock_guardrails/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_guardrails/lambda/src/app.py
@@ -1,4 +1,4 @@
-"""This script performs operations to enable, configure, and disable Bedrock security controls.
+"""This script performs operations to create, configure, and delete Bedrock guardrails.
 
 Version: 1.0
 
@@ -28,6 +28,7 @@
 import sra_repo
 import sra_s3
 import sra_sns
+import sra_sqs
 import sra_ssm_params
 import sra_sts
 
@@ -53,7 +54,6 @@ def load_kms_key_policies() -> dict:
 SOLUTION_NAME: str = "sra-bedrock-guardrails"
 GOVERNED_REGIONS = []
 ORGANIZATION_ID = ""
-# SRA_ALARM_EMAIL: str = ""
 SRA_ALARM_TOPIC_ARN: str = ""
 STATE_TABLE: str = "sra_state"  # for saving resource info
 CFN_CUSTOM_RESOURCE: str = "Custom::LambdaCustomResource"
@@ -173,6 +173,7 @@ def load_kms_key_policies() -> dict:
 cloudwatch = sra_cloudwatch.SRACloudWatch()
 kms = sra_kms.SRAKMS()
 bedrock = sra_bedrock.SRABedrock()
+sqs = sra_sqs.SRASQS()
 
 # propagate solution name to class objects
 cloudwatch.SOLUTION_NAME = SOLUTION_NAME
@@ -191,7 +192,6 @@ def get_resource_parameters(event: dict) -> None:
     global DRY_RUN
     global GOVERNED_REGIONS
     global CFN_RESPONSE_DATA
-    # global SRA_ALARM_EMAIL
     global ORGANIZATION_ID
 
     param_validation: dict = validate_parameters(event["ResourceProperties"], PARAMETER_VALIDATION_RULES)
@@ -238,9 +238,6 @@ def get_resource_parameters(event: dict) -> None:
         LOGGER.info("Error retrieving SRA staging bucket ssm parameter.  Is the SRA common prerequisites solution deployed?")
         raise ValueError("Error retrieving SRA staging bucket ssm parameter.  Is the SRA common prerequisites solution deployed?") from None
 
-    # if event["ResourceProperties"]["SRA_ALARM_EMAIL"] != "":
-    #     SRA_ALARM_EMAIL = event["ResourceProperties"]["SRA_ALARM_EMAIL"]
-
     if event["ResourceProperties"]["DRY_RUN"] == "true":
         # dry run
         LOGGER.info("Dry run enabled...")
@@ -586,6 +583,33 @@ def create_kms_key(acct: str, region: str) -> None:
             )
 
 
+def check_sqs_queue() -> str:
+    """Add sqs queue record if DLQ exists.
+
+    Returns:
+        str: sns topic arn
+    """
+    global DRY_RUN_DATA
+    global LIVE_RUN_DATA
+    global CFN_RESPONSE_DATA
+
+    sns.SNS_CLIENT = sts.assume_role(sts.MANAGEMENT_ACCOUNT, sts.CONFIGURATION_ROLE, "sns", sts.HOME_REGION)
+    queue_search = sqs.find_sqs_queue(f"{SOLUTION_NAME}-DLQ")
+    if queue_search is None:
+        LOGGER.info(f"{SOLUTION_NAME}-DLQ doesn't exist")
+
+    else:
+        LOGGER.info(f"{SOLUTION_NAME}-DLQ sqs queue exists.")
+        queue_arn = queue_search
+        if DRY_RUN is False:
+            # SQS State table record:
+            add_state_table_record("sqs", "implemented", "sqs queue", "queue", queue_arn, ACCOUNT, sts.HOME_REGION, f"{SOLUTION_NAME}-DLQ")
+        else:
+            DRY_RUN_DATA["SQSCreate"] = f"DRY_RUN: {SOLUTION_NAME}-DLQ sqs queue exists"
+
+    return queue_arn
+
+
 def create_guardrail(acct: str, region: str, params: dict) -> None:
     """Deploy the Bedrock guardrail.
 
@@ -631,7 +655,6 @@ def set_guardrail_config(params: dict, guardrail_key_id: str) -> Dict:
         "description": "sra bedrock guardrail",
         "blockedInputMessaging": params["BLOCKED_INPUT_MESSAGING"],
         "blockedOutputsMessaging": params["BLOCKED_OUTPUTS_MESSAGING"],
-        # "clientRequestToken": 'sra-client-request-token-12',
         "kmsKeyId": guardrail_key_id,
         "tags": [
             {"key": "sra-solution", "value": params["SOLUTION_NAME"]},
@@ -668,6 +691,75 @@ def set_guardrail_config(params: dict, guardrail_key_id: str) -> Dict:
     return guardrail_params
 
 
+def delete_bedrock_guardrails_key(acct: str, region: str) -> None:
+    """Delete KMS key.
+
+    Args:
+        acct (str): AWS account ID
+        region (str): AWS region name
+    """
+    # Delete KMS key (schedule deletion) and delete kms alias
+    kms.KMS_CLIENT = sts.assume_role(acct, sts.CONFIGURATION_ROLE, "kms", region)
+    search_bedrock_guardrails_kms_key, bedrock_guardrails_key_alias, bedrock_guardrails_key_id, bedrock_guardrails_key_arn = kms.check_alias_exists(
+        kms.KMS_CLIENT, f"alias/{GUARDRAILS_KEY_ALIAS}"
+    )
+    if search_bedrock_guardrails_kms_key is True:
+        if DRY_RUN is False:
+            LOGGER.info(f"Deleting {GUARDRAILS_KEY_ALIAS} KMS key")
+            kms.delete_alias(kms.KMS_CLIENT, f"alias/{GUARDRAILS_KEY_ALIAS}")
+            LIVE_RUN_DATA[f"KMSDeleteAlias-{acct}-{region}"] = f"Deleted {GUARDRAILS_KEY_ALIAS} KMS key alias"
+            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
+            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
+            LOGGER.info(f"Deleting {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
+            remove_state_table_record(bedrock_guardrails_key_arn)
+
+            kms.schedule_key_deletion(kms.KMS_CLIENT, bedrock_guardrails_key_id)
+            LIVE_RUN_DATA[f"KMSDelete-{acct}-{region}"] = f"Deleted {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})"
+            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
+            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
+            LOGGER.info(f"Scheduled deletion of {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
+            kms_key_arn = f"arn:{sts.PARTITION}:kms:{region}:{acct}:key/{bedrock_guardrails_key_id}"
+            remove_state_table_record(kms_key_arn)
+
+        else:
+            LOGGER.info(f"DRY_RUN: Deleting {GUARDRAILS_KEY_ALIAS} KMS key")
+            DRY_RUN_DATA[f"KMSAliasDelete-{acct}-{region}"] = f"DRY_RUN: Delete {GUARDRAILS_KEY_ALIAS} KMS key"
+            LOGGER.info(f"DRY_RUN: Deleting {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
+            DRY_RUN_DATA[f"KMSDelete-{acct}-{region}"] = f"DRY_RUN: Delete {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})"
+    else:
+        LOGGER.info(f"{GUARDRAILS_KEY_ALIAS} KMS key does not exist.")
+
+
+def delete_guardrails(account: str, region: str, guardrail_name: str) -> None:
+    """Delete the Bedrock guardrails.
+
+    Args:
+        account: AWS account id
+        region: AWS region
+        guardrail_name: Name of the Bedrock guardrail to delete.
+    """
+    global DRY_RUN_DATA
+    global LIVE_RUN_DATA
+    global CFN_RESPONSE_DATA
+
+    if DRY_RUN is False:
+        bedrock.BEDROCK_CLIENT = sts.assume_role(account, sts.CONFIGURATION_ROLE, "bedrock", region)
+        LOGGER.info(f"Deleting Bedrock guardrail in {account} in {region}...")
+        guardrail_id = bedrock.get_guardrail_id(guardrail_name)
+        if guardrail_id != "":
+            bedrock.delete_guardrail(guardrail_id)
+            LIVE_RUN_DATA[f"Bedrock-guardrail-{account}_{region}"] = f"Deleted Bedrock Guardrail ({guardrail_name}) in {account} in {region}"
+            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
+            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
+            guardrail_arn = f"arn:aws:bedrock:{region}:{account}:guardrail/{guardrail_id}"
+            remove_state_table_record(guardrail_arn)
+        else:
+            LOGGER.info(f"Guardrail {guardrail_name} does not exist in {account} in {region}")
+    else:
+        LOGGER.info(f"DRY_RUN: Delete Bedrock guardrail {guardrail_name} in {account} in {region}")
+        DRY_RUN_DATA[f"Bedrock-guardrail-{account}_{region}"] = f"DRY_RUN: Delete Bedrock guardrail {guardrail_name}"
+
+
 def create_event(event: dict, context: Any) -> str:
     """Create event.
 
@@ -731,7 +823,7 @@ def create_event(event: dict, context: Any) -> str:
                 LOGGER.info(f"Guardrail {event['ResourceProperties']['BEDROCK_GUARDRAIL_NAME']} does not exist in {acct} in {region}")
                 create_kms_key(acct, region)
                 create_guardrail(acct, region, event["ResourceProperties"])
-
+    check_sqs_queue()
     # End
     if DRY_RUN is False:
         LOGGER.info(json.dumps({"RUN STATS": CFN_RESPONSE_DATA, "RUN DATA": LIVE_RUN_DATA}))
@@ -771,75 +863,6 @@ def update_event(event: dict, context: Any) -> str:
     return CFN_RESOURCE_ID
 
 
-def delete_bedrock_guardrails_key(acct: str, region: str) -> None:
-    """Delete KMS key.
-
-    Args:
-        acct (str): AWS account ID
-        region (str): AWS region name
-    """
-    # Delete KMS key (schedule deletion) and delete kms alias
-    kms.KMS_CLIENT = sts.assume_role(acct, sts.CONFIGURATION_ROLE, "kms", region)
-    search_bedrock_guardrails_kms_key, bedrock_guardrails_key_alias, bedrock_guardrails_key_id, bedrock_guardrails_key_arn = kms.check_alias_exists(
-        kms.KMS_CLIENT, f"alias/{GUARDRAILS_KEY_ALIAS}"
-    )
-    if search_bedrock_guardrails_kms_key is True:
-        if DRY_RUN is False:
-            LOGGER.info(f"Deleting {GUARDRAILS_KEY_ALIAS} KMS key")
-            kms.delete_alias(kms.KMS_CLIENT, f"alias/{GUARDRAILS_KEY_ALIAS}")
-            LIVE_RUN_DATA[f"KMSDeleteAlias-{acct}-{region}"] = f"Deleted {GUARDRAILS_KEY_ALIAS} KMS key alias"
-            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
-            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
-            LOGGER.info(f"Deleting {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
-            remove_state_table_record(bedrock_guardrails_key_arn)
-
-            kms.schedule_key_deletion(kms.KMS_CLIENT, bedrock_guardrails_key_id)
-            LIVE_RUN_DATA[f"KMSDelete-{acct}-{region}"] = f"Deleted {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})"
-            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
-            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
-            LOGGER.info(f"Scheduled deletion of {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
-            kms_key_arn = f"arn:{sts.PARTITION}:kms:{region}:{acct}:key/{bedrock_guardrails_key_id}"
-            remove_state_table_record(kms_key_arn)
-
-        else:
-            LOGGER.info(f"DRY_RUN: Deleting {GUARDRAILS_KEY_ALIAS} KMS key")
-            DRY_RUN_DATA[f"KMSAliasDelete-{acct}-{region}"] = f"DRY_RUN: Delete {GUARDRAILS_KEY_ALIAS} KMS key"
-            LOGGER.info(f"DRY_RUN: Deleting {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})")
-            DRY_RUN_DATA[f"KMSDelete-{acct}-{region}"] = f"DRY_RUN: Delete {GUARDRAILS_KEY_ALIAS} KMS key ({bedrock_guardrails_key_id})"
-    else:
-        LOGGER.info(f"{GUARDRAILS_KEY_ALIAS} KMS key does not exist.")
-
-
-def delete_guardrails(account: str, region: str, guardrail_name: str) -> None:
-    """Delete the Bedrock guardrails.
-
-    Args:
-        account: AWS account id
-        region: AWS region
-        guardrail_name: Name of the Bedrock guardrail to delete.
-    """
-    global DRY_RUN_DATA
-    global LIVE_RUN_DATA
-    global CFN_RESPONSE_DATA
-
-    if DRY_RUN is False:
-        bedrock.BEDROCK_CLIENT = sts.assume_role(account, sts.CONFIGURATION_ROLE, "bedrock", region)
-        LOGGER.info(f"Deleting Bedrock guardrail in {account} in {region}...")
-        guardrail_id = bedrock.get_guardrail_id(guardrail_name)
-        if guardrail_id != "":
-            bedrock.delete_guardrail(guardrail_id)
-            LIVE_RUN_DATA[f"Bedrock-guardrail-{account}_{region}"] = f"Deleted Bedrock Guardrail ({guardrail_name}) in {account} in {region}"
-            CFN_RESPONSE_DATA["deployment_info"]["action_count"] += 1
-            CFN_RESPONSE_DATA["deployment_info"]["resources_deployed"] -= 1
-            guardrail_arn = f"arn:aws:bedrock:{region}:{account}:guardrail/{guardrail_id}"
-            remove_state_table_record(guardrail_arn)
-        else:
-            LOGGER.info(f"Guardrail {guardrail_name} does not exist in {account} in {region}")
-    else:
-        LOGGER.info(f"DRY_RUN: Delete Bedrock guardrail {guardrail_name} in {account} in {region}")
-        DRY_RUN_DATA[f"Bedrock-guardrail-{account}_{region}"] = f"DRY_RUN: Delete Bedrock guardrail {guardrail_name}"
-
-
 def delete_event(event: dict, context: Any) -> None:  # noqa: CFQ001, CCR001, C901
     """Delete event function.
 
@@ -854,15 +877,18 @@ def delete_event(event: dict, context: Any) -> None:  # noqa: CFQ001, CCR001, C9
     LIVE_RUN_DATA = {}
     LOGGER.info("Delete event function")
 
-    # 4) Delete Bedrock guardrails
+    # Delete Bedrock guardrails
     accounts, regions = get_accounts_and_regions(
         event["ResourceProperties"]["SRA_BEDROCK_ACCOUNTS"], event["ResourceProperties"]["SRA_BEDROCK_REGIONS"]
     )
     for acct in accounts:
         for region in regions:
             delete_guardrails(acct, region, event["ResourceProperties"]["BEDROCK_GUARDRAIL_NAME"])
             delete_bedrock_guardrails_key(acct, region)
-
+    # Remove sqs queue record
+    queue_arn = check_sqs_queue()
+    if queue_arn is not None:
+        remove_state_table_record(queue_arn)
     # 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}")
diff --git a/aws_sra_examples/solutions/genai/bedrock_guardrails/lambda/src/sra_sqs.py b/aws_sra_examples/solutions/genai/bedrock_guardrails/lambda/src/sra_sqs.py
@@ -0,0 +1,77 @@
+"""Lambda module to setup SRA SQS resources in the organization.
+
+Version: 0.1
+
+SQS module 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
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+from time import sleep
+from typing import TYPE_CHECKING
+
+import boto3
+import sra_sts
+from botocore.config import Config
+from botocore.exceptions import ClientError
+
+if TYPE_CHECKING:
+    from mypy_boto3_sqs.client import SQSClient
+
+
+class SRASQS:
+    """Class to setup SRA SQS resources in the organization."""
+
+    # Setup Default Logger
+    LOGGER = logging.getLogger(__name__)
+    log_level: str = os.environ.get("LOG_LEVEL", "INFO")
+    LOGGER.setLevel(log_level)
+
+    BOTO3_CONFIG = Config(retries={"max_attempts": 10, "mode": "standard"})
+    UNEXPECTED = "Unexpected!"
+
+    try:
+        MANAGEMENT_ACCOUNT_SESSION = boto3.Session()
+        SQS_CLIENT: SQSClient = MANAGEMENT_ACCOUNT_SESSION.client("sqs", config=BOTO3_CONFIG)
+    except Exception:
+        LOGGER.exception(UNEXPECTED)
+        raise ValueError("Unexpected error executing Lambda function. Review CloudWatch logs for details.") from None
+
+    sts = sra_sts.SRASTS()
+
+    def find_sqs_queue(self, queue_name: str, region: str = "default", account: str = "default") -> str | None:
+        """Find SQS Queue ARN.
+
+        Args:
+            queue_name (str): SQS Queue Name
+            region (str): AWS Region
+            account (str): AWS Account
+
+        Raises:
+            ValueError: Error finding SQS Queue
+
+        Returns:
+            str: SQS Queue ARN
+        """
+        if region == "default":
+            region = self.sts.HOME_REGION
+        if account == "default":
+            account = self.sts.MANAGEMENT_ACCOUNT
+        try:
+            response = self.SQS_CLIENT.get_queue_attributes(
+                QueueUrl=f"https://sqs.{region}.amazonaws.com/{account}/{queue_name}", AttributeNames=["QueueArn"]
+            )
+            return response["Attributes"]["QueueArn"]
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NotFoundException":
+                self.LOGGER.info(f"SQS Queue '{queue_name}' not found exception.")
+                return None
+            if e.response["Error"]["Code"] == "NotFound":
+                self.LOGGER.info(f"SQS Queue '{queue_name}' not found.")
+                return None
+            raise ValueError(f"Error finding SQS topic: {e}") from None
diff --git a/aws_sra_examples/solutions/genai/bedrock_guardrails/templates/sra-bedrock-guardrails-main.yaml b/aws_sra_examples/solutions/genai/bedrock_guardrails/templates/sra-bedrock-guardrails-main.yaml
@@ -1,5 +1,6 @@
 AWSTemplateFormatVersion: '2010-09-09'
-Description: CloudFormation template to create a Lambda function and its execution role
+Description: CloudFormation template to create a Lambda function and its execution role to deploy Bedrock guardrails - 'bedrock_guardrails' solution in the repo,
+  https://github.com/aws-samples/aws-security-reference-architecture-examples(sra-1u3sd7f8i)
 
 Metadata:
   AWS::CloudFormation::Interface:
@@ -224,7 +225,6 @@ Parameters:
       '^\[(?:\s*\{\s*\"type\"\s*:\s*\"(ADDRESS|AGE|AWS_ACCESS_KEY|AWS_SECRET_KEY|CA_HEALTH_NUMBER|CA_SOCIAL_INSURANCE_NUMBER|CREDIT_DEBIT_CARD_CVV|CREDIT_DEBIT_CARD_EXPIRY|CREDIT_DEBIT_CARD_NUMBER|DRIVER_ID|EMAIL|INTERNATIONAL_BANK_ACCOUNT_NUMBER|IP_ADDRESS|LICENSE_PLATE|MAC_ADDRESS|NAME|PASSWORD|PHONE|PIN|SWIFT_CODE|UK_NATIONAL_HEALTH_SERVICE_NUMBER|UK_NATIONAL_INSURANCE_NUMBER|UK_UNIQUE_TAXPAYER_REFERENCE_NUMBER|URL|USERNAME|US_BANK_ACCOUNT_NUMBER|US_BANK_ROUTING_NUMBER|US_INDIVIDUAL_TAX_IDENTIFICATION_NUMBER|US_PASSPORT_NUMBER|US_SOCIAL_SECURITY_NUMBER|VEHICLE_IDENTIFICATION_NUMBER)\"\s*,\s*\"action\"\s*:\s*\"(BLOCK|ANONYMIZE)\"\s*\})(?!.*\"type\"\s*:\s*\"\1\")(?:\s*,\s*\{\s*\"type\"\s*:\s*\"(ADDRESS|AGE|AWS_ACCESS_KEY|AWS_SECRET_KEY|CA_HEALTH_NUMBER|CA_SOCIAL_INSURANCE_NUMBER|CREDIT_DEBIT_CARD_CVV|CREDIT_DEBIT_CARD_EXPIRY|CREDIT_DEBIT_CARD_NUMBER|DRIVER_ID|EMAIL|INTERNATIONAL_BANK_ACCOUNT_NUMBER|IP_ADDRESS|LICENSE_PLATE|MAC_ADDRESS|NAME|PASSWORD|PHONE|PIN|SWIFT_CODE|UK_NATIONAL_HEALTH_SERVICE_NUMBER|UK_NATIONAL_INSURANCE_NUMBER|UK_UNIQUE_TAXPAYER_REFERENCE_NUMBER|URL|USERNAME|US_BANK_ACCOUNT_NUMBER|US_BANK_ROUTING_NUMBER|US_INDIVIDUAL_TAX_IDENTIFICATION_NUMBER|US_PASSPORT_NUMBER|US_SOCIAL_SECURITY_NUMBER|VEHICLE_IDENTIFICATION_NUMBER)\"\s*,\s*\"action\"\s*:\s*\"(BLOCK|ANONYMIZE)\"\s*\}(?!.*\"type\"\s*:\s*\"\3\"))*\s*\]$'
     ConstraintDescription: Must be a valid JSON array of objects. Each object must have a 'type' (valid PII entity type) and an 'action' (BLOCK or ANONYMIZE). At least one configuration is required.
 
-  # consider adding GuardrailRegexConfig
   pDeployContextualGroundingPolicy:
     Type: String
     Default: 'true'
@@ -296,7 +296,8 @@ Resources:
             Statement:
               - Effect: Allow
                 Action:
-                  - sqs:SendMessage
+                  - 'sqs:SendMessage'
+                  - 'sqs:GetQueueAttributes'
                 Resource: !GetAtt rBedrockGuardrailsDLQ.Arn
         - PolicyDocument:
             Version: '2012-10-17'

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# SRA Bedrock Organizations Solution`
	`1`	`+# SRA Bedrock Guardrails Solution`
`2`	`2`
`3`	`3`	`## Table of Contents`
`4`	`4`	`- [Introduction](#introduction)`