Add support for aws govcloud in dist.sh (#510)

.internal/aws/cloudformation/macro.yaml

@@ -121,7 +121,7 @@ Resources:
             parameters = event["templateParameterValues"]
             if "ElasticServerlessForwarderS3ConfigFile" in parameters:
                 bucket_name_and_object_key = parameters["ElasticServerlessForwarderS3ConfigFile"].replace("s3://", "")
-                resource = f"arn:aws:s3:::{bucket_name_and_object_key}"
+                resource = f"arn:%awsOrGov%:s3:::{bucket_name_and_object_key}"
                 if len(resource) > 0:
                     policy_fragment["Properties"]["PolicyDocument"]["Statement"].append(
                         {
@@ -246,7 +246,7 @@ Resources:
   MacroElasticServerlessForwarder:
     Type: AWS::CloudFormation::Macro
     Properties:
-      Description: Expand parameters to Events and Policy for %sarAppName%
+      Description: Expand parameters to Events and Policy for %sarAppName% in %awsOrGov%
       FunctionName: !GetAtt MacroElasticServerlessForwarderFunction.Arn
       Name: %sarAppName%-macro
 Metadata:

.internal/aws/cloudformation/template.yaml

@@ -51,13 +51,13 @@ Resources:
     Type: AWS::Serverless::Application
     Properties:
       Location:
-        ApplicationId: arn:aws:serverlessrepo:%awsRegion%:%accountID%:applications/helper-macro-%sarAppName%
+        ApplicationId: arn:%awsOrGov%:serverlessrepo:%awsRegion%:%accountID%:applications/helper-macro-%sarAppName%
         SemanticVersion: %semanticVersion%
   ElasticServerlessForwarderApplication:
     Type: AWS::Serverless::Application
     Properties:
       Location:
-        ApplicationId: arn:aws:serverlessrepo:%awsRegion%:%accountID%:applications/helper-application-%sarAppName%
+        ApplicationId: arn:%awsOrGov%:serverlessrepo:%awsRegion%:%accountID%:applications/helper-application-%sarAppName%
         SemanticVersion: %semanticVersion%
       Parameters:
         ElasticServerlessForwarderS3ConfigFile: !Ref ElasticServerlessForwarderS3ConfigFile

.internal/aws/scripts/dist.sh

@@ -35,7 +35,18 @@ CODE_URI="${TMPDIR}/sources"
 
 trap "rm -rf ${TMPDIR}" EXIT
 
-aws s3api get-bucket-location --bucket "${BUCKET}" || aws s3api create-bucket --acl private --bucket "${BUCKET}" --region "${REGION}" --create-bucket-configuration LocationConstraint="${REGION}"
+aws s3api get-bucket-location --bucket "${BUCKET}" --region "${REGION}" || aws s3api create-bucket --acl private --bucket "${BUCKET}" --region "${REGION}" --create-bucket-configuration LocationConstraint="${REGION}"
+
+# Check if region is in AWS GovCloud and create bucket arn
+if [[ ${REGION} == *"$gov"* ]]; then
+  BUCKET_ARN="arn:aws-us-gov:s3:::${BUCKET}"
+  AWS_OR_AWS_GOV="aws-us-gov"
+else
+  BUCKET_ARN="arn:aws:s3:::${BUCKET}"
+  AWS_OR_AWS_GOV="aws"
+fi
+
+BUCKET_RESOURCE="${BUCKET_ARN}/*"
 
 cat <<EOF > "${TMPDIR}/policy.json"
 {
@@ -47,7 +58,7 @@ cat <<EOF > "${TMPDIR}/policy.json"
                 "Service":  "serverlessrepo.amazonaws.com"
             },
             "Action": "s3:GetObject",
-            "Resource": "arn:aws:s3:::${BUCKET}/*",
+            "Resource": "${BUCKET_RESOURCE}",
             "Condition" : {
                 "StringEquals": {
                     "aws:SourceAccount": "${ACCOUNT_ID}"
@@ -58,7 +69,7 @@ cat <<EOF > "${TMPDIR}/policy.json"
 }
 EOF
 
-aws s3api put-bucket-policy --bucket "${BUCKET}" --policy "file://${TMPDIR}/policy.json"
+aws s3api put-bucket-policy --bucket "${BUCKET}" --region "${REGION}" --policy "file://${TMPDIR}/policy.json"
 mkdir -v -p "${CODE_URI}"
 cp -v requirements.txt "${CODE_URI}/"
 cp -v main_aws.py "${CODE_URI}/"
@@ -67,8 +78,8 @@ find {handlers,share,shippers,storage} -not -name "*__pycache__*" -name "*.py" -
 cp -v LICENSE.txt "${CODE_URI}/LICENSE.txt"
 cp -v docs/README-AWS.md "${CODE_URI}/README.md"
 
-sed -e "s|%codeUri%|${CODE_URI}|g" -e "s/%sarAppName%/${SAR_APP_NAME}/g" -e "s/%sarAuthorName%/${SAR_AUTHOR_NAME}/g" -e "s/%semanticVersion%/${SEMANTIC_VERSION}/g" -e "s/%awsRegion%/${REGION}/g" .internal/aws/cloudformation/macro.yaml > "${TMPDIR}/macro.yaml"
-sed -e "s|%codeUri%|${CODE_URI}|g" -e "s/%sarAppName%/${SAR_APP_NAME}/g" -e "s/%sarAuthorName%/${SAR_AUTHOR_NAME}/g" -e "s/%semanticVersion%/${SEMANTIC_VERSION}/g" -e "s/%awsRegion%/${REGION}/g" -e "s/%accountID%/${ACCOUNT_ID}/g" .internal/aws/cloudformation/template.yaml > "${TMPDIR}/template.yaml"
+sed -e "s|%codeUri%|${CODE_URI}|g" -e "s/%sarAppName%/${SAR_APP_NAME}/g" -e "s/%sarAuthorName%/${SAR_AUTHOR_NAME}/g" -e "s/%semanticVersion%/${SEMANTIC_VERSION}/g" -e "s/%awsRegion%/${REGION}/g" -e "s/%awsOrGov%/${AWS_OR_AWS_GOV}/g" .internal/aws/cloudformation/macro.yaml > "${TMPDIR}/macro.yaml"
+sed -e "s|%codeUri%|${CODE_URI}|g" -e "s/%sarAppName%/${SAR_APP_NAME}/g" -e "s/%sarAuthorName%/${SAR_AUTHOR_NAME}/g" -e "s/%semanticVersion%/${SEMANTIC_VERSION}/g" -e "s/%awsRegion%/${REGION}/g" -e "s/%accountID%/${ACCOUNT_ID}/g" -e "s/%awsOrGov%/${AWS_OR_AWS_GOV}/g" .internal/aws/cloudformation/template.yaml > "${TMPDIR}/template.yaml"
 sed -e "s|%codeUri%|${CODE_URI}|g" -e "s/%sarAppName%/${SAR_APP_NAME}/g" -e "s/%sarAuthorName%/${SAR_AUTHOR_NAME}/g" -e "s/%semanticVersion%/${SEMANTIC_VERSION}/g" -e "s/%awsRegion%/${REGION}/g" -e "s/%codeURIBucket%/${BUCKET}/g" .internal/aws/cloudformation/application.yaml > "${TMPDIR}/application.yaml"
 
 sam build --debug --use-container --build-dir "${TMPDIR}/.aws-sam/build/macro" --template-file "${TMPDIR}/macro.yaml" --region "${REGION}"

handlers/aws/utils.py

@@ -268,10 +268,13 @@ def get_continuing_original_input_type(sqs_record: dict[str, Any]) -> Optional[s
 
     original_event_source: str = sqs_record["messageAttributes"]["originalEventSourceARN"]["stringValue"]
 
-    if original_event_source.startswith("arn:aws:logs"):
+    if original_event_source.startswith("arn:aws:logs") or original_event_source.startswith("arn:aws-us-gov:logs"):
         return "cloudwatch-logs"
 
-    if original_event_source.startswith("arn:aws:kinesis") and original_event_source.find(":stream/") > -1:
+    if (
+        original_event_source.startswith("arn:aws:kinesis")
+        or original_event_source.startswith("arn:aws-us-gov:kinesis")
+    ) and original_event_source.find(":stream/") > -1:
         return "kinesis-data-stream"
 
     return None
@@ -392,7 +395,14 @@ def get_input_from_log_group_subscription_data(
     assert "Regions" in all_regions
     for region_data in all_regions["Regions"]:
         region = region_data["RegionName"]
-        log_stream_arn = f"arn:aws:logs:{region}:{account_id}:log-group:{log_group_name}:log-stream:{log_stream_name}"
+
+        aws_or_gov = "aws"
+        if "gov" in region:
+            aws_or_gov = "aws-us-gov"
+
+        log_stream_arn = (
+            f"arn:{aws_or_gov}:logs:{region}:{account_id}:log-group:{log_group_name}:log-stream:{log_stream_name}"
+        )
         event_input = config.get_input_by_id(log_stream_arn)
 
         if event_input is not None:

publish_lambda.sh

@@ -39,7 +39,18 @@ GIT_REPO="https://github.com/elastic/elastic-serverless-forwarder.git"
 
 trap 'rm -rf ${TMPDIR}' EXIT
 
-aws s3api get-bucket-location --bucket "${BUCKET}" || aws s3api create-bucket --acl private --bucket "${BUCKET}" --region "${REGION}" --create-bucket-configuration LocationConstraint="${REGION}"
+aws s3api get-bucket-location --bucket "${BUCKET}" --region "${REGION}" || aws s3api create-bucket --acl private --bucket "${BUCKET}" --region "${REGION}" --create-bucket-configuration LocationConstraint="${REGION}"
+
+# Check if region is in AWS GovCloud and create bucket arn
+if [[ ${REGION} == *"$gov"* ]]; then
+  BUCKET_ARN="arn:aws-us-gov:s3:::${BUCKET}"
+  AWS_OR_AWS_GOV="aws-us-gov"
+else
+  BUCKET_ARN="arn:aws:s3:::${BUCKET}"
+  AWS_OR_AWS_GOV="aws"
+fi
+
+BUCKET_RESOURCE="${BUCKET_ARN}/*"
 
 mkdir -v -p "${CLONED_FOLDER}"
 git clone --depth 1 --branch "${TAG_NAME}" "${GIT_REPO}" "${CLONED_FOLDER}"
@@ -322,7 +333,7 @@ def create_policy(publish_config: dict[str, Any]):
         assert isinstance(publish_config["s3-config-file"], str)
 
         bucket_name_and_object_key = publish_config["s3-config-file"].replace("s3://", "")
-        resource = f"arn:aws:s3:::{bucket_name_and_object_key}"
+        resource = f"arn:${AWS_OR_AWS_GOV}:s3:::{bucket_name_and_object_key}"
         if len(resource) > 0:
             policy_fragment["Properties"]["PolicyDocument"]["Statement"].append(
                 {"Effect": "Allow", "Action": "s3:GetObject", "Resource": resource}
@@ -471,18 +482,18 @@ if __name__ == "__main__":
             },
             "RoleName": "${CUSTOM_ROLE_PREFIX}ApplicationElasticServerlessForwarderRole",
             "ManagedPolicyArns": [
-                "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
-                "arn:aws:iam::aws:policy/service-role/AWSLambdaSQSQueueExecutionRole"
+                "arn:${AWS_OR_AWS_GOV}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
+                "arn:${AWS_OR_AWS_GOV}:iam::aws:policy/service-role/AWSLambdaSQSQueueExecutionRole"
             ],
         },
     }
 
     if vpc_config:
-        customRole["Properties"]["ManagedPolicyArns"].append("arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole")
+        customRole["Properties"]["ManagedPolicyArns"].append("arn:${AWS_OR_AWS_GOV}:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole")
 
     has_kinesis_events: bool = len([created_event for created_event in created_events if created_events[created_event]["Type"] == "Kinesis"]) > 0
     if has_kinesis_events:
-        customRole["Properties"]["ManagedPolicyArns"].append("arn:aws:iam::aws:policy/service-role/AWSLambdaKinesisExecutionRole")
+        customRole["Properties"]["ManagedPolicyArns"].append("arn:${AWS_OR_AWS_GOV}:iam::aws:policy/service-role/AWSLambdaKinesisExecutionRole")
 
     cloudformation_yaml["Resources"]["ApplicationElasticServerlessForwarderCustomRole"] = customRole
     cloudformation_yaml["Resources"]["ApplicationElasticServerlessForwarder"]["Properties"]["Role"] = {"Fn::GetAtt": ["ApplicationElasticServerlessForwarderCustomRole", "Arn"] }