[AWS] fix(logs-forwarder): properly set AWS partition for log groups tags cache (#940)

* [AWS] fix(logs-forwarder): properly set AWS partition for log groups tags cache

* fix template

Author: ge0Aja

aws/logs_monitoring/steps/handlers/aws_attributes.py

@@ -2,15 +2,26 @@
 
 
 class AwsAttributes:
-    def __init__(self, log_group=None, log_stream=None, log_events=None, owner=None):
+    def __init__(
+        self, context, log_group=None, log_stream=None, log_events=None, owner=None
+    ):
         self.log_group = log_group
         self.log_stream = log_stream
         self.log_events = log_events
         self.owner = owner
+        self.partition = self._get_aws_partition(context)
         self.lambda_arn = None
         self.account = None
         self.region = None
 
+    def _get_aws_partition(self, context):
+        if context.invoked_function_arn.startswith("arn:aws-cn:"):
+            return "aws-cn"
+        elif context.invoked_function_arn.startswith("arn:aws-us-gov:"):
+            return "aws-us-gov"
+        else:
+            return "aws"
+
     def to_dict(self):
         awslogs = {
             "aws": {
@@ -30,7 +41,7 @@ def get_log_group(self):
         return self.log_group
 
     def get_log_group_arn(self):
-        return f"arn:aws:logs:{self.region}:{self.account}:log-group:{self.log_group}"
+        return f"arn:{self.partition}:logs:{self.region}:{self.account}:log-group:{self.log_group}"
 
     def get_log_stream(self):
         return self.log_stream

aws/logs_monitoring/steps/handlers/awslogs_handler.py

@@ -37,6 +37,7 @@ def handle(self, event):
         logs = self.extract_logs(event)
         # Build aws attributes
         aws_attributes = AwsAttributes(
+            self.context,
             logs.get("logGroup"),
             logs.get("logStream"),
             logs.get("logEvents"),

aws/logs_monitoring/tests/test_awslogs_handler.py

@@ -33,6 +33,9 @@
 
 
 class Context:
+    def __init__(self, invoked_function_arn="invoked_function_arn"):
+        self.invoked_function_arn = invoked_function_arn
+
     function_version = "$LATEST"
     invoked_function_arn = "invoked_function_arn"
     function_name = "function_name"
@@ -45,6 +48,7 @@ def setUp(self):
             r"forwarder_version:\d+\.\d+\.\d+",
             "forwarder_version:<redacted>",
         )
+        self.context = Context()
 
     @patch("caching.cloudwatch_log_group_cache.CloudwatchLogGroupTagsCache.__init__")
     def test_handle_with_overridden_source(self, mock_cache_init):
@@ -79,13 +83,12 @@ def test_handle_with_overridden_source(self, mock_cache_init):
             }
 
             # Create required args
-            context = Context()
             mock_cache_init.return_value = None
             cache_layer = CacheLayer("")
             cache_layer._cloudwatch_log_group_cache.get = MagicMock(return_value=[])
 
             # Process the event
-            awslogs_handler = AwsLogsHandler(context, cache_layer)
+            awslogs_handler = AwsLogsHandler(self.context, cache_layer)
 
             # Verify
             verify_as_json(
@@ -127,14 +130,13 @@ def test_awslogs_handler_rds_postgresql(self, mock_cache_init):
                 )
             }
         }
-        context = Context()
         mock_cache_init.return_value = None
         cache_layer = CacheLayer("")
         cache_layer._cloudwatch_log_group_cache.get = MagicMock(
             return_value=["test_tag_key:test_tag_value"]
         )
 
-        awslogs_handler = AwsLogsHandler(context, cache_layer)
+        awslogs_handler = AwsLogsHandler(self.context, cache_layer)
         verify_as_json(
             list(awslogs_handler.handle(event)),
             options=Options().with_scrubber(self.scrubber),
@@ -189,7 +191,6 @@ def test_awslogs_handler_step_functions_tags_added_properly(
                 )
             }
         }
-        context = Context()
         mock_forward_metrics.side_effect = MagicMock()
         mock_cache_init.return_value = None
         cache_layer = CacheLayer("")
@@ -198,7 +199,7 @@ def test_awslogs_handler_step_functions_tags_added_properly(
         )
         cache_layer._cloudwatch_log_group_cache.get = MagicMock()
 
-        awslogs_handler = AwsLogsHandler(context, cache_layer)
+        awslogs_handler = AwsLogsHandler(self.context, cache_layer)
         verify_as_json(
             list(awslogs_handler.handle(event)),
             options=Options().with_scrubber(self.scrubber),
@@ -252,7 +253,6 @@ def test_awslogs_handler_step_functions_customized_log_group(
                 )
             }
         }
-        context = Context()
         mock_forward_metrics.side_effect = MagicMock()
         mock_cache_init.return_value = None
         cache_layer = CacheLayer("")
@@ -261,7 +261,7 @@ def test_awslogs_handler_step_functions_customized_log_group(
         )
         cache_layer._cloudwatch_log_group_cache.get = MagicMock()
 
-        awslogs_handler = AwsLogsHandler(context, cache_layer)
+        awslogs_handler = AwsLogsHandler(self.context, cache_layer)
         # for some reasons, the below two are needed to update the context of the handler
         verify_as_json(
             list(awslogs_handler.handle(eventFromCustomizedLogGroup)),
@@ -304,14 +304,13 @@ def test_awslogs_handler_lambda_log(self):
                 )
             }
         }
-        context = Context()
         cache_layer = CacheLayer("")
         cache_layer._cloudwatch_log_group_cache.get = MagicMock()
         cache_layer._lambda_cache.get = MagicMock(
             return_value=["service:customtags_service"]
         )
 
-        awslogs_handler = AwsLogsHandler(context, cache_layer)
+        awslogs_handler = AwsLogsHandler(self.context, cache_layer)
         verify_as_json(
             list(awslogs_handler.handle(event)),
             options=Options().with_scrubber(self.scrubber),
@@ -327,12 +326,12 @@ def test_process_lambda_logs(self):
         }
         metadata = {"ddsource": "postgresql", "ddtags": ""}
         aws_attributes = AwsAttributes(
+            self.context,
             stepfunction_loggroup.get("logGroup"),
             stepfunction_loggroup.get("logStream"),
             stepfunction_loggroup.get("owner"),
         )
-        context = Context()
-        aws_handler = AwsLogsHandler(context, CacheLayer(""))
+        aws_handler = AwsLogsHandler(self.context, CacheLayer(""))
 
         aws_handler.process_lambda_logs(metadata, aws_attributes)
         self.assertEqual(metadata, {"ddsource": "postgresql", "ddtags": ""})
@@ -346,13 +345,13 @@ def test_process_lambda_logs(self):
         }
         metadata = {"ddsource": "postgresql", "ddtags": "env:dev"}
         aws_attributes = AwsAttributes(
+            self.context,
             lambda_default_loggroup.get("logGroup"),
             lambda_default_loggroup.get("logStream"),
             lambda_default_loggroup.get("owner"),
         )
-        context = Context()
 
-        aws_handler = AwsLogsHandler(context, CacheLayer(""))
+        aws_handler = AwsLogsHandler(self.context, CacheLayer(""))
         aws_handler.process_lambda_logs(metadata, aws_attributes)
         self.assertEqual(
             metadata,
@@ -380,12 +379,13 @@ def test_process_lambda_logs(self):
 
 class TestLambdaCustomizedLogGroup(unittest.TestCase):
     def setUp(self):
-        self.aws_handler = AwsLogsHandler(None, None)
+        self.context = Context()
+        self.aws_handler = AwsLogsHandler(self.context, None)
 
     def test_get_lower_cased_lambda_function_name(self):
-        self.assertEqual(True, True)
         # Non Lambda log
         aws_attributes = AwsAttributes(
+            self.context,
             "/aws/vendedlogs/states/logs-to-traces-sequential-Logs",
             "states/logs-to-traces-sequential/2022-11-10-15-50/7851b2d9",
             [],
@@ -396,6 +396,7 @@ def test_get_lower_cased_lambda_function_name(self):
         )
 
         aws_attributes = AwsAttributes(
+            self.context,
             "/aws/lambda/test-lambda-default-log-group",
             "2023/11/06/[$LATEST]b25b1f977b3e416faa45a00f427e7acb",
             [],
@@ -406,6 +407,7 @@ def test_get_lower_cased_lambda_function_name(self):
         )
 
         aws_attributes = AwsAttributes(
+            self.context,
             "customizeLambdaGrop",
             "2023/11/06/test-customized-log-group1[$LATEST]13e304cba4b9446eb7ef082a00038990",
             [],
@@ -418,20 +420,23 @@ def test_get_lower_cased_lambda_function_name(self):
 
 class TestParsingStepFunctionLogs(unittest.TestCase):
     def setUp(self):
-        self.aws_handler = AwsLogsHandler(None, None)
+        self.context = Context()
+        self.aws_handler = AwsLogsHandler(self.context, None)
 
     def test_get_state_machine_arn(self):
         aws_attributes = AwsAttributes(
+            context=self.context,
             log_events=[
                 {
                     "message": json.dumps({"no_execution_arn": "xxxx/yyy"}),
                 }
-            ]
+            ],
         )
 
         self.assertEqual(self.aws_handler.get_state_machine_arn(aws_attributes), "")
 
         aws_attributes = AwsAttributes(
+            context=self.context,
             log_events=[
                 {
                     "message": json.dumps(
@@ -442,14 +447,15 @@ def test_get_state_machine_arn(self):
                         }
                     ),
                 }
-            ]
+            ],
         )
         self.assertEqual(
             self.aws_handler.get_state_machine_arn(aws_attributes),
             "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States",
         )
 
         aws_attributes = AwsAttributes(
+            context=self.context,
             log_events=[
                 {
                     "message": json.dumps(
@@ -460,7 +466,7 @@ def test_get_state_machine_arn(self):
                         }
                     )
                 }
-            ]
+            ],
         )
 
         self.assertEqual(
@@ -469,6 +475,7 @@ def test_get_state_machine_arn(self):
         )
 
         aws_attributes = AwsAttributes(
+            context=self.context,
             log_events=[
                 {
                     "message": json.dumps(
@@ -479,13 +486,66 @@ def test_get_state_machine_arn(self):
                         }
                     )
                 }
-            ]
+            ],
         )
         self.assertEqual(
             self.aws_handler.get_state_machine_arn(aws_attributes),
             "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States",
         )
 
 
+class TestAwsPartitionExtraction(unittest.TestCase):
+    def test_get_log_group_aws_partition(self):
+        # default partition
+        context = Context(
+            invoked_function_arn="arn:aws:lambda:us-east-1:12345678910:function:test-lambda"
+        )
+        aws_attributes = AwsAttributes(
+            context=context,
+            log_group="my-log-group",
+        )
+
+        aws_attributes.set_account_region(
+            "arn:aws:lambda:us-east-1:12345678910:function:test-lambda"
+        )
+
+        self.assertEqual(
+            aws_attributes.get_log_group_arn(),
+            "arn:aws:logs:us-east-1:12345678910:log-group:my-log-group",
+        )
+
+        # aws-cn partition
+        context = Context(
+            invoked_function_arn="arn:aws-cn:lambda:cn-north-1:12345678910:function:test-lambda"
+        )
+        aws_attributes = AwsAttributes(
+            context=context,
+            log_group="my-log-group",
+        )
+        aws_attributes.set_account_region(
+            "arn:aws-cn:lambda:cn-north-1:12345678910:function:test-lambda"
+        )
+        self.assertEqual(
+            aws_attributes.get_log_group_arn(),
+            "arn:aws-cn:logs:cn-north-1:12345678910:log-group:my-log-group",
+        )
+
+        # aws-us-gov partition
+        context = Context(
+            invoked_function_arn="arn:aws-us-gov:lambda:us-gov-west-1:12345678910:function:test-lambda"
+        )
+        aws_attributes = AwsAttributes(
+            context=context,
+            log_group="my-log-group",
+        )
+        aws_attributes.set_account_region(
+            "arn:aws-us-gov:lambda:us-gov-west-1:12345678910:function:test-lambda"
+        )
+        self.assertEqual(
+            aws_attributes.get_log_group_arn(),
+            "arn:aws-us-gov:logs:us-gov-west-1:12345678910:log-group:my-log-group",
+        )
+
+
 if __name__ == "__main__":
     unittest.main()