CLOUDS-2607:Properly parse partitions for gov and china service arns (#598)

* CLOUDS-2607:Properly parse partitions for gov and china service arns

* Handle redshift gov events

* Fix gov partition

* Formatting

Author: saleelsaptarshi-dd

aws/logs_monitoring/parsing.py

@@ -36,6 +36,8 @@
     DD_USE_VPC,
 )
 
+GOV, CN = "gov", "cn"
+
 
 logger = logging.getLogger()
 
@@ -378,6 +380,16 @@ def find_s3_source(key):
     return "s3"
 
 
+def get_partition_from_region(region):
+    partition = "aws"
+    if region:
+        if GOV in region:
+            partition = "aws-us-gov"
+        elif CN in region:
+            partition = "aws-cn"
+    return partition
+
+
 def parse_service_arn(source, key, bucket, context):
     if source == "elb":
         # For ELB logs we parse the filename to extract parameters in order to rebuild the ARN
@@ -406,8 +418,9 @@ def parse_service_arn(source, key, bucket, context):
             elbname = name.replace(".", "/")
             if len(idsplit) > 1:
                 idvalue = idsplit[1]
-                return "arn:aws:elasticloadbalancing:{}:{}:loadbalancer/{}".format(
-                    region, idvalue, elbname
+                partition = get_partition_from_region(region)
+                return "arn:{}:elasticloadbalancing:{}:{}:loadbalancer/{}".format(
+                    partition, region, idvalue, elbname
                 )
     if source == "s3":
         # For S3 access logs we use the bucket name to rebuild the arn
@@ -446,8 +459,8 @@ def parse_service_arn(source, key, bucket, context):
             filesplit = filename.split("_")
             if len(filesplit) == 6:
                 clustername = filesplit[3]
-                return "arn:aws:redshift:{}:{}:cluster:{}:".format(
-                    region, accountID, clustername
+                return "arn:{}:redshift:{}:{}:cluster:{}:".format(
+                    get_partition_from_region(region), region, accountID, clustername
                 )
     return
 

aws/logs_monitoring/tests/test_parsing.py

@@ -213,6 +213,16 @@ def test_redshift_event(self):
             "redshift",
         )
 
+    def test_redshift_gov_event(self):
+        self.assertEqual(
+            parse_event_source(
+                {"Records": ["logs-from-s3"]},
+                "AWSLogs/123456779121/redshift/us-gov-east-1/2020/10/21/123456779121_redshift_us-gov-east"
+                "-1_mycluster_userlog_2020-10-21T18:01.gz",
+            ),
+            "redshift",
+        )
+
     def test_route53_event(self):
         self.assertEqual(
             parse_event_source(
@@ -358,6 +368,20 @@ def test_elb_s3_key_multi_prefix(self):
             "arn:aws:elasticloadbalancing:us-east-1:123456789123:loadbalancer/app/my-alb-name/123456789aabcdef",
         )
 
+    def test_elb_s3_key_multi_prefix_gov(self):
+        self.assertEqual(
+            parse_service_arn(
+                "elb",
+                "elasticloadbalancing/my-alb-name/AWSLogs/123456789123/elasticloadbalancing/us-gov-east-1/2022/02/08"
+                "/123456789123_elasticloadbalancing_us-gov-east-1_app.my-alb-name.123456789aabcdef_20220208T1127Z_10"
+                ".0.0.2_1abcdef2.log.gz",
+                None,
+                None,
+            ),
+            "arn:aws-us-gov:elasticloadbalancing:us-gov-east-1:123456789123:loadbalancer/app/my-alb-name"
+            "/123456789aabcdef",
+        )
+
 
 class TestParseAwsWafLogs(unittest.TestCase):
     def test_waf_string_invalid_json(self):