s3-sqs-image-pipeline-kms

import json
import os
import logging
from datetime import datetime, timezone
from urllib.parse import unquote_plus

import boto3

log = logging.getLogger()
log.setLevel(logging.INFO)

DEST_BUCKET = os.environ["DEST_BUCKET"]
DEST_KMS_KEY_ID = os.environ["DEST_KMS_KEY_ID"]
DDB_TABLE = os.environ["DDB_TABLE"]

s3 = boto3.client("s3")
ddb = boto3.client("dynamodb")


def _process_s3_record(s3_rec):
    src_bucket = s3_rec["s3"]["bucket"]["name"]
    src_key = unquote_plus(s3_rec["s3"]["object"]["key"])

    obj = s3.get_object(Bucket=src_bucket, Key=src_key)
    body = obj["Body"].read()
    size = len(body)

    dest_key = f"processed/{src_key}"
    s3.put_object(
        Bucket=DEST_BUCKET,
        Key=dest_key,
        Body=body,
        ServerSideEncryption="aws:kms",
        SSEKMSKeyId=DEST_KMS_KEY_ID,
    )

    ddb.put_item(
        TableName=DDB_TABLE,
        Item={
            "object_key": {"S": src_key},
            "source_size": {"N": str(size)},
            "processed_at": {"S": datetime.now(timezone.utc).isoformat()},
            "status": {"S": "OK"},
        },
    )


def lambda_handler(event, context):
    failures = []
    for record in event.get("Records", []):
        message_id = record["messageId"]
        try:
            body = json.loads(record["body"])
            # Direct S3 event notifications have Records at top level.
            # Skip S3 test events (contain "Event": "s3:TestEvent").
            if body.get("Event") == "s3:TestEvent":
                continue
            s3_records = body.get("Records", [])
            for s3_rec in s3_records:
                _process_s3_record(s3_rec)
        except Exception as e:
            log.exception("failed processing message %s: %s", message_id, e)
            failures.append({"itemIdentifier": message_id})

    return {"batchItemFailures": failures}

#!/usr/bin/env bash
set -euo pipefail

export AWS_ACCESS_KEY_ID=test
export AWS_SECRET_ACCESS_KEY=test
export AWS_DEFAULT_REGION=us-east-1

EP=http://localstack:4566
AWS="aws --endpoint-url=$EP"
ACCOUNT=000000000000
REGION=us-east-1

SRC_BUCKET=ingest-source-bucket
DST_BUCKET=ingest-dest-bucket
MAIN_QUEUE=ingest-events
DLQ=ingest-dlq
TABLE=image-ingest-records
FN=image-ingest-consumer
ROLE=image-ingest-lambda-role

ROLE_ARN="arn:aws:iam::${ACCOUNT}:role/${ROLE}"

echo "== creating KMS keys =="

# Source bucket CMK: root admin + lambda role decrypt
cat >/tmp/src-key-policy.json <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "RootAdmin",
      "Effect": "Allow",
      "Principal": {"AWS": "arn:aws:iam::${ACCOUNT}:root"},
      "Action": "kms:*",
      "Resource": "*"
    },
    {
      "Sid": "LambdaDecrypt",
      "Effect": "Allow",
      "Principal": {"AWS": "${ROLE_ARN}"},
      "Action": ["kms:Decrypt", "kms:DescribeKey"],
      "Resource": "*"
    }
  ]
}
EOF

SRC_KMS_ID=$($AWS kms create-key --description "source bucket cmk" \
  --policy file:///tmp/src-key-policy.json --query KeyMetadata.KeyId --output text)
$AWS kms create-alias --alias-name alias/source-bucket-cmk --target-key-id $SRC_KMS_ID
SRC_KMS_ARN="arn:aws:kms:${REGION}:${ACCOUNT}:key/${SRC_KMS_ID}"
echo "src kms: $SRC_KMS_ARN"

# SQS CMK: root admin + lambda decrypt + s3 service to encrypt
cat >/tmp/sqs-key-policy.json <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "RootAdmin",
      "Effect": "Allow",
      "Principal": {"AWS": "arn:aws:iam::${ACCOUNT}:root"},
      "Action": "kms:*",
      "Resource": "*"
    },
    {
      "Sid": "S3Publisher",
      "Effect": "Allow",
      "Principal": {"Service": "s3.amazonaws.com"},
      "Action": ["kms:GenerateDataKey", "kms:Decrypt"],
      "Resource": "*"
    },
    {
      "Sid": "LambdaConsumer",
      "Effect": "Allow",
      "Principal": {"AWS": "${ROLE_ARN}"},
      "Action": ["kms:Decrypt", "kms:GenerateDataKey", "kms:DescribeKey"],
      "Resource": "*"
    }
  ]
}
EOF

SQS_KMS_ID=$($AWS kms create-key --description "sqs queue cmk" \
  --policy file:///tmp/sqs-key-policy.json --query KeyMetadata.KeyId --output text)
$AWS kms create-alias --alias-name alias/sqs-queue-cmk --target-key-id $SQS_KMS_ID
SQS_KMS_ARN="arn:aws:kms:${REGION}:${ACCOUNT}:key/${SQS_KMS_ID}"
echo "sqs kms: $SQS_KMS_ARN"

# Dest bucket CMK: root admin + lambda encrypt/generate
cat >/tmp/dst-key-policy.json <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "RootAdmin",
      "Effect": "Allow",
      "Principal": {"AWS": "arn:aws:iam::${ACCOUNT}:root"},
      "Action": "kms:*",
      "Resource": "*"
    },
    {
      "Sid": "LambdaEncrypt",
      "Effect": "Allow",
      "Principal": {"AWS": "${ROLE_ARN}"},
      "Action": ["kms:Encrypt", "kms:GenerateDataKey", "kms:DescribeKey"],
      "Resource": "*"
    }
  ]
}
EOF

DST_KMS_ID=$($AWS kms create-key --description "dest bucket cmk" \
  --policy file:///tmp/dst-key-policy.json --query KeyMetadata.KeyId --output text)
$AWS kms create-alias --alias-name alias/dest-bucket-cmk --target-key-id $DST_KMS_ID
DST_KMS_ARN="arn:aws:kms:${REGION}:${ACCOUNT}:key/${DST_KMS_ID}"
echo "dst kms: $DST_KMS_ARN"

# persist for later steps
cat >/tmp/vars.env <<EOF
SRC_KMS_ID=$SRC_KMS_ID
SRC_KMS_ARN=$SRC_KMS_ARN
SQS_KMS_ID=$SQS_KMS_ID
SQS_KMS_ARN=$SQS_KMS_ARN
DST_KMS_ID=$DST_KMS_ID
DST_KMS_ARN=$DST_KMS_ARN
EOF

echo "== creating DynamoDB table =="
$AWS dynamodb create-table \
  --table-name $TABLE \
  --attribute-definitions AttributeName=object_key,AttributeType=S \
  --key-schema AttributeName=object_key,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST >/dev/null

echo "== creating buckets =="
$AWS s3api create-bucket --bucket $SRC_BUCKET >/dev/null
$AWS s3api create-bucket --bucket $DST_BUCKET >/dev/null

# SSE-KMS default encryption on both buckets
cat >/tmp/src-enc.json <<EOF
{"Rules":[{"ApplyServerSideEncryptionByDefault":{"SSEAlgorithm":"aws:kms","KMSMasterKeyID":"$SRC_KMS_ARN"},"BucketKeyEnabled":true}]}
EOF
$AWS s3api put-bucket-encryption --bucket $SRC_BUCKET \
  --server-side-encryption-configuration file:///tmp/src-enc.json

cat >/tmp/dst-enc.json <<EOF
{"Rules":[{"ApplyServerSideEncryptionByDefault":{"SSEAlgorithm":"aws:kms","KMSMasterKeyID":"$DST_KMS_ARN"},"BucketKeyEnabled":true}]}
EOF
$AWS s3api put-bucket-encryption --bucket $DST_BUCKET \
  --server-side-encryption-configuration file:///tmp/dst-enc.json

echo "setup phase 1 done"

#!/usr/bin/env bash
set -euo pipefail

export AWS_ACCESS_KEY_ID=test
export AWS_SECRET_ACCESS_KEY=test
export AWS_DEFAULT_REGION=us-east-1

EP=http://localstack:4566
AWS="aws --endpoint-url=$EP"
ACCOUNT=000000000000
REGION=us-east-1

SRC_BUCKET=ingest-source-bucket
DST_BUCKET=ingest-dest-bucket
MAIN_QUEUE=ingest-events
DLQ=ingest-dlq
TABLE=image-ingest-records
FN=image-ingest-consumer
ROLE=image-ingest-lambda-role

# load vars
source /tmp/vars.env

ROLE_ARN="arn:aws:iam::${ACCOUNT}:role/${ROLE}"
MAIN_ARN="arn:aws:sqs:${REGION}:${ACCOUNT}:${MAIN_QUEUE}"
DLQ_ARN="arn:aws:sqs:${REGION}:${ACCOUNT}:${DLQ}"
SRC_ARN="arn:aws:s3:::${SRC_BUCKET}"
DST_ARN="arn:aws:s3:::${DST_BUCKET}"
TABLE_ARN="arn:aws:dynamodb:${REGION}:${ACCOUNT}:table/${TABLE}"

echo "== creating DLQ =="
$AWS sqs create-queue --queue-name $DLQ >/dev/null

echo "== creating main queue with SSE-KMS + redrive =="
REDRIVE=$(printf '{"deadLetterTargetArn":"%s","maxReceiveCount":"3"}' "$DLQ_ARN")
$AWS sqs create-queue --queue-name $MAIN_QUEUE \
  --attributes "{\"VisibilityTimeout\":\"90\",\"KmsMasterKeyId\":\"$SQS_KMS_ARN\",\"RedrivePolicy\":$(echo $REDRIVE | jq -Rs .)}" \
  >/dev/null

# queue policy allowing s3 source bucket to SendMessage
cat >/tmp/queue-policy.json <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AllowS3Publish",
      "Effect": "Allow",
      "Principal": {"Service": "s3.amazonaws.com"},
      "Action": "sqs:SendMessage",
      "Resource": "${MAIN_ARN}",
      "Condition": {
        "ArnLike": {"aws:SourceArn": "${SRC_ARN}"}
      }
    }
  ]
}
EOF

POLICY_JSON=$(cat /tmp/queue-policy.json | jq -Rs .)
$AWS sqs set-queue-attributes --queue-url ${EP}/000000000000/${MAIN_QUEUE} \
  --attributes "{\"Policy\":${POLICY_JSON}}"

echo "== creating IAM role =="
cat >/tmp/trust.json <<'EOF'
{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"lambda.amazonaws.com"},"Action":"sts:AssumeRole"}]}
EOF
$AWS iam create-role --role-name $ROLE --assume-role-policy-document file:///tmp/trust.json >/dev/null

cat >/tmp/role-policy.json <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "Logs",
      "Effect": "Allow",
      "Action": ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"],
      "Resource": "arn:aws:logs:${REGION}:${ACCOUNT}:*"
    },
    {
      "Sid": "SQSConsume",
      "Effect": "Allow",
      "Action": [
        "sqs:ReceiveMessage",
        "sqs:DeleteMessage",
        "sqs:GetQueueAttributes",
        "sqs:ChangeMessageVisibility"
      ],
      "Resource": "${MAIN_ARN}"
    },
    {
      "Sid": "S3Read",
      "Effect": "Allow",
      "Action": "s3:GetObject",
      "Resource": "${SRC_ARN}/*"
    },
    {
      "Sid": "S3Write",
      "Effect": "Allow",
      "Action": "s3:PutObject",
      "Resource": "${DST_ARN}/*"
    },
    {
      "Sid": "DdbWrite",
      "Effect": "Allow",
      "Action": "dynamodb:PutItem",
      "Resource": "${TABLE_ARN}"
    },
    {
      "Sid": "KMSSourceDecrypt",
      "Effect": "Allow",
      "Action": ["kms:Decrypt", "kms:DescribeKey"],
      "Resource": "${SRC_KMS_ARN}"
    },
    {
      "Sid": "KMSSqsDecrypt",
      "Effect": "Allow",
      "Action": ["kms:Decrypt", "kms:GenerateDataKey", "kms:DescribeKey"],
      "Resource": "${SQS_KMS_ARN}"
    },
    {
      "Sid": "KMSDestEncrypt",
      "Effect": "Allow",
      "Action": ["kms:Encrypt", "kms:GenerateDataKey", "kms:DescribeKey"],
      "Resource": "${DST_KMS_ARN}"
    }
  ]
}
EOF
$AWS iam put-role-policy --role-name $ROLE --policy-name ${ROLE}-policy \
  --policy-document file:///tmp/role-policy.json

echo "== packaging lambda =="
cd /app
cp handler.py /tmp/handler.py
(cd /tmp && zip -q handler.zip handler.py)

echo "== creating lambda =="
$AWS lambda create-function \
  --function-name $FN \
  --runtime python3.11 \
  --role $ROLE_ARN \
  --handler handler.lambda_handler \
  --timeout 30 \
  --zip-file fileb:///tmp/handler.zip \
  --environment "Variables={DEST_BUCKET=${DST_BUCKET},DEST_KMS_KEY_ID=${DST_KMS_ARN},DDB_TABLE=${TABLE}}" \
  >/dev/null

# wait for active
for i in $(seq 1 30); do
  state=$($AWS lambda get-function --function-name $FN --query Configuration.State --output text 2>/dev/null || echo pending)
  [ "$state" = "Active" ] && break
  sleep 1
done
echo "lambda state: $state"

echo "== event source mapping SQS -> Lambda with partial batch failures =="
$AWS lambda create-event-source-mapping \
  --function-name $FN \
  --event-source-arn $MAIN_ARN \
  --batch-size 10 \
  --function-response-types ReportBatchItemFailures \
  >/dev/null

echo "== S3 bucket notification -> SQS =="
cat >/tmp/notif.json <<EOF
{
  "QueueConfigurations": [
    {
      "Id": "ingest-created",
      "QueueArn": "${MAIN_ARN}",
      "Events": ["s3:ObjectCreated:*"]
    }
  ]
}
EOF
$AWS s3api put-bucket-notification-configuration --bucket $SRC_BUCKET \
  --notification-configuration file:///tmp/notif.json

echo "setup phase 2 done"