diff --git a/infrastructure/terraform/components/notifyai/lambda.tf b/infrastructure/terraform/components/notifyai/lambda.tf index b25e78d..b403a86 100644 --- a/infrastructure/terraform/components/notifyai/lambda.tf +++ b/infrastructure/terraform/components/notifyai/lambda.tf @@ -108,6 +108,7 @@ resource "aws_lambda_function" "bedrock-messager" { env_logging_s3_key_prefix = local.s3_lambda_logging_key env_guardrail_arn = aws_bedrock_guardrail.notifai-bedrock-guardrail.guardrail_arn env_guardrail_version = "DRAFT" + env_logging_s3_account_id = var.aws_account_id } } } diff --git a/src/backend/bedrock-prompt-messager/core/config.py b/src/backend/bedrock-prompt-messager/core/config.py index 98dde26..16845c4 100644 --- a/src/backend/bedrock-prompt-messager/core/config.py +++ b/src/backend/bedrock-prompt-messager/core/config.py @@ -12,3 +12,4 @@ def __init__(self): self.logging_s3_key_prefix = os.environ.get("env_logging_s3_key_prefix") self.guardrail = os.environ.get("env_guardrail_arn") self.guardrail_version = os.environ.get("env_guardrail_version") + self.logging_s3_account_id = os.environ.get("env_logging_s3_account_id") diff --git a/src/backend/bedrock-prompt-messager/services/bedrock_service.py b/src/backend/bedrock-prompt-messager/services/bedrock_service.py index 412f25c..7b6ce3c 100644 --- a/src/backend/bedrock-prompt-messager/services/bedrock_service.py +++ b/src/backend/bedrock-prompt-messager/services/bedrock_service.py @@ -135,7 +135,7 @@ def _get_admail_tool_config(self): def log_prompt_details_to_s3( self, promptinput, promptoutput, guardrail_assessment, filename ): - if not self.config.logging_s3_bucket or not self.config.logging_s3_key_prefix: + if not self.config.logging_s3_bucket or not self.config.logging_s3_key_prefix or not self.config.logging_s3_account_id: print(constants.ERROR_S3_LOGGING_NOT_CONFIGURED) return @@ -164,6 +164,7 @@ def log_prompt_details_to_s3( Key=s3_key, Body=json.dumps(log_data, indent=4), ContentType="application/json", + ExpectedBucketOwner=self.config.logging_s3_account_id, ) except Exception as e: print(f"Error logging to S3: {e}") diff --git a/src/backend/bedrock-prompt-messager/tests/conftest.py b/src/backend/bedrock-prompt-messager/tests/conftest.py index a55da7c..cec5217 100644 --- a/src/backend/bedrock-prompt-messager/tests/conftest.py +++ b/src/backend/bedrock-prompt-messager/tests/conftest.py @@ -12,3 +12,4 @@ def set_env_vars(monkeypatch): monkeypatch.setenv("env_logging_s3_key_prefix", "logs/") monkeypatch.setenv("env_guardrail_arn", "test-arn") monkeypatch.setenv("env_guardrail_version", "1") + monkeypatch.setenv("env_logging_s3_account_id", "123456789012")