fix: Allow SNS config has separate list of regions

This fix enables the SNS config to specify a
list of regions for sending notifications.

Users can send notifications to specific regions,
but the notification will be sent in every regions
if not specified.

Author: JessicaJang

awspub/configmodels.py

@@ -112,6 +112,12 @@ class ConfigImageSNSNotificationModel(BaseModel):
         description="The body of the message to be sent to subscribers.",
         default={SNSNotificationProtocol.DEFAULT: ""},
     )
+    regions: Optional[List[str]] = Field(
+        description="Optional list of regions for sending notification. If not given, regions where the image "
+        "registered will be used from the currently used parition. If a region doesn't exist in the currently "
+        "used partition, it will be ignored.",
+        default=None,
+    )
 
     @field_validator("message")
     def check_message(cls, value):

awspub/image.py

@@ -342,15 +342,7 @@ def _sns_publish(self) -> None:
         Publish SNS notifiations about newly available images to subscribers
         """
 
-        # Checking if the image(s) are registered or published before sending the notification
-        for region in self.image_regions:
-            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
-            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
-            if not image_info:
-                logger.error(f"can not send SNS notification for {self.image_name} because no image found in {region}")
-                return
-
-        SNSNotification(self._ctx, self.image_name, self._s3.bucket_region).publish()
+        SNSNotification(self._ctx, self.image_name).publish()
 
     def cleanup(self) -> None:
         """

awspub/sns.py

@@ -40,7 +40,6 @@ def conf(self) -> List[Dict[str, Any]]:
         """
         return self._ctx.conf["images"][self._image_name]["sns"]
 
-    def _get_topic_arn(self, topic_name: str) -> str:
     def _sns_regions(self, topic_config: Dict[Any, Any]) -> List[str]:
         """
         Get the sns regions. Either configured in the sns configuration
@@ -55,6 +54,7 @@ def _sns_regions(self, topic_config: Dict[Any, Any]) -> List[str]:
 
         return sns_regions
 
+    def _get_topic_arn(self, topic_name: str, region_name: str) -> str:
         """
         Calculate topic ARN based on partition, region, account and topic name
         :param topic_name: Name of topic
@@ -65,40 +65,41 @@ def _sns_regions(self, topic_config: Dict[Any, Any]) -> List[str]:
         :rtype: str
         """
 
-        stsclient: STSClient = boto3.client("sts", region_name=self._region_name)
+        stsclient: STSClient = boto3.client("sts", region_name=region_name)
         resp = stsclient.get_caller_identity()
 
         account = resp["Account"]
         # resp["Arn"] has string format "arn:partition:iam::accountnumber:user/iam_role"
         partition = resp["Arn"].rsplit(":")[1]
 
-        return f"arn:{partition}:sns:{self._region_name}:{account}:{topic_name}"
+        return f"arn:{partition}:sns:{region_name}:{account}:{topic_name}"
 
     def publish(self) -> None:
         """
         send notification to subscribers
         """
 
-        snsclient: SNSClient = boto3.client("sns", region_name=self._region_name)
-
         for topic in self.conf:
             for topic_name, topic_config in topic.items():
-                try:
-                    snsclient.publish(
-                        TopicArn=self._get_topic_arn(topic_name),
-                        Subject=topic_config["subject"],
-                        Message=json.dumps(topic_config["message"]),
-                        MessageStructure="json",
-                    )
-                except ClientError as e:
-                    exception_code: str = e.response["Error"]["Code"]
-                    if exception_code == "AuthorizationError":
-                        raise AWSAuthorizationException(
-                            "Profile does not have a permission to send the SNS notification. Please review the policy."
+                for region_name in self._sns_regions(topic_config):
+                    snsclient: SNSClient = boto3.client("sns", region_name=region_name)
+                    try:
+                        snsclient.publish(
+                            TopicArn=self._get_topic_arn(topic_name, region_name),
+                            Subject=topic_config["subject"],
+                            Message=json.dumps(topic_config["message"]),
+                            MessageStructure="json",
                         )
-                    else:
-                        raise AWSNotificationException(str(e))
-                logger.info(
-                    f"The SNS notification {topic_config['subject']}"
-                    f" for the topic {topic_name} in {self._region_name} has been sent."
-                )
+                    except ClientError as e:
+                        exception_code: str = e.response["Error"]["Code"]
+                        if exception_code == "AuthorizationError":
+                            raise AWSAuthorizationException(
+                                "Profile does not have a permission to send the SNS notification."
+                                " Please review the policy."
+                            )
+                        else:
+                            raise AWSNotificationException(str(e))
+                    logger.info(
+                        f"The SNS notification {topic_config['subject']}"
+                        f" for the topic {topic_name} in {region_name} has been sent."
+                    )

awspub/tests/fixtures/config1.yaml

@@ -130,6 +130,8 @@ awspub:
             message:
               default: "default-message"
               email: "email-message"
+            regions:
+              - "us-east-1"
     "test-image-11":
       boot_mode: "uefi"
       description: |
@@ -143,10 +145,27 @@ awspub:
             message:
               default: "default-message"
               email: "email-message"
+            regions:
+              - "us-east-1"
         - "topic2":
             subject: "topic2-subject"
             message:
               default: "default-message"
+            regions:
+              - "us-gov-1"
+              - "eu-central-1"
+    "test-image-12":
+      boot_mode: "uefi"
+      description: |
+        A test image without a separate snapshot but single sns configs
+      regions:
+        - "us-east-1"
+      sns:
+        - "topic1":
+            subject: "topic1-subject"
+            message:
+              default: "default-message"
+              email: "email-message"
 
   tags:
     name: "foobar"

awspub/tests/test_api.py

@@ -25,6 +25,7 @@
                 "test-image-9",
                 "test-image-10",
                 "test-image-11",
+                "test-image-12",
             ],
         ),
         # with a group that no image as, no image should be processed

awspub/tests/test_image.py

@@ -143,6 +143,7 @@ def test_image___get_root_device_snapshot_id(root_device_name, block_device_mapp
         ("test-image-8", "aws-cn", True, True, False, True, False),
         ("test-image-10", "aws", False, False, False, False, True),
         ("test-image-11", "aws", False, False, False, False, True),
+        ("test-image-12", "aws", False, False, False, False, True),
     ],
 )
 def test_image_publish(
@@ -183,7 +184,6 @@ def test_image_publish(
             "Regions": [{"RegionName": "eu-central-1"}, {"RegionName": "us-east-1"}]
         }
         instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
-        instance.list_topics.return_value = {"Topics": [{"TopicArn": "arn:aws:sns:topic1"}]}
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
         img = image.Image(ctx, imagename)
         img.publish()

awspub/tests/test_sns.py

@@ -13,7 +13,8 @@
     "imagename,called_sns_publish, publish_call_count",
     [
         ("test-image-10", True, 1),
-        ("test-image-11", True, 4),
+        ("test-image-11", True, 2),
+        ("test-image-12", True, 2),
     ],
 )
 def test_sns_publish(imagename, called_sns_publish, publish_call_count):
@@ -23,11 +24,12 @@ def test_sns_publish(imagename, called_sns_publish, publish_call_count):
     with patch("boto3.client") as bclient_mock:
         instance = bclient_mock.return_value
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
-        image_conf = ctx.conf["images"][imagename]
-
-        for region in image_conf["regions"]:
-            sns.SNSNotification(ctx, imagename, region).publish()
+        instance.describe_regions.return_value = {
+            "Regions": [{"RegionName": "eu-central-1"}, {"RegionName": "us-east-1"}]
+        }
+        instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
 
+        sns.SNSNotification(ctx, imagename).publish()
         assert instance.publish.called == called_sns_publish
         assert instance.publish.call_count == publish_call_count
 
@@ -37,6 +39,7 @@ def test_sns_publish(imagename, called_sns_publish, publish_call_count):
     [
         ("test-image-10"),
         ("test-image-11"),
+        ("test-image-12"),
     ],
 )
 def test_sns_publish_fail_with_invalid_topic(imagename):
@@ -46,12 +49,15 @@ def test_sns_publish_fail_with_invalid_topic(imagename):
     with patch("boto3.client") as bclient_mock:
         instance = bclient_mock.return_value
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
-        image_conf = ctx.conf["images"][imagename]
+        instance.describe_regions.return_value = {
+            "Regions": [{"RegionName": "eu-central-1"}, {"RegionName": "us-east-1"}]
+        }
+        instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
 
         # topic1 is invalid topic
         def side_effect(*args, **kwargs):
             topic_arn = kwargs.get("TopicArn")
-            if "topic1" in topic_arn:
+            if "topic1" in topic_arn and "us-east-1" in topic_arn:
                 error_reponse = {
                     "Error": {
                         "Code": "NotFoundException",
@@ -63,16 +69,16 @@ def side_effect(*args, **kwargs):
 
         instance.publish.side_effect = side_effect
 
-        for region in image_conf["regions"]:
-            with pytest.raises(exceptions.AWSNotificationException):
-                sns.SNSNotification(ctx, imagename, region).publish()
+        with pytest.raises(exceptions.AWSNotificationException):
+            sns.SNSNotification(ctx, imagename).publish()
 
 
 @pytest.mark.parametrize(
     "imagename",
     [
         ("test-image-10"),
         ("test-image-11"),
+        ("test-image-12"),
     ],
 )
 def test_sns_publish_fail_with_unauthorized_user(imagename):
@@ -82,7 +88,10 @@ def test_sns_publish_fail_with_unauthorized_user(imagename):
     with patch("boto3.client") as bclient_mock:
         instance = bclient_mock.return_value
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
-        image_conf = ctx.conf["images"][imagename]
+        instance.describe_regions.return_value = {
+            "Regions": [{"RegionName": "eu-central-1"}, {"RegionName": "us-east-1"}]
+        }
+        instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
 
         error_reponse = {
             "Error": {
@@ -92,49 +101,89 @@ def test_sns_publish_fail_with_unauthorized_user(imagename):
         }
         instance.publish.side_effect = botocore.exceptions.ClientError(error_reponse, "")
 
-        for region in image_conf["regions"]:
-            with pytest.raises(exceptions.AWSAuthorizationException):
-                sns.SNSNotification(ctx, imagename, region).publish()
+        with pytest.raises(exceptions.AWSAuthorizationException):
+            sns.SNSNotification(ctx, imagename).publish()
 
 
 @pytest.mark.parametrize(
-    "imagename, partition, expected",
+    "imagename, partition, regions_in_partition, expected",
     [
         (
             "test-image-10",
             "aws-cn",
+            ["cn-north1", "cn-northwest-1"],
+            [],
+        ),
+        (
+            "test-image-11",
+            "aws",
+            ["us-east-1", "eu-central-1"],
             [
-                "arn:aws-cn:sns:us-east-1:1234:topic1",
+                "arn:aws:sns:us-east-1:1234:topic1",
+                "arn:aws:sns:eu-central-1:1234:topic2",
             ],
         ),
         (
-            "test-image-11",
+            "test-image-12",
             "aws",
+            ["us-east-1", "eu-central-1"],
             [
                 "arn:aws:sns:us-east-1:1234:topic1",
-                "arn:aws:sns:us-east-1:1234:topic2",
                 "arn:aws:sns:eu-central-1:1234:topic1",
-                "arn:aws:sns:eu-central-1:1234:topic2",
             ],
         ),
     ],
 )
-def test_sns__get_topic_arn(imagename, partition, expected):
+def test_sns__get_topic_arn(imagename, partition, regions_in_partition, expected):
     """
     Test the send_notification logic
     """
     with patch("boto3.client") as bclient_mock:
         instance = bclient_mock.return_value
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
-        image_conf = ctx.conf["images"][imagename]
+        sns_conf = ctx.conf["images"][imagename]["sns"]
+        instance.describe_regions.return_value = {"Regions": [{"RegionName": r} for r in regions_in_partition]}
+        instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
 
         instance.get_caller_identity.return_value = {"Account": "1234", "Arn": f"arn:{partition}:iam::1234:user/test"}
 
         topic_arns = []
-        for region in image_conf["regions"]:
-            for topic in image_conf["sns"]:
-                topic_name = next(iter(topic))
-                res_arn = sns.SNSNotification(ctx, imagename, region)._get_topic_arn(topic_name)
-                topic_arns.append(res_arn)
+        for topic in sns_conf:
+            for topic_name, topic_conf in topic.items():
+                sns_regions = sns.SNSNotification(ctx, imagename)._sns_regions(topic_conf)
+                for region in sns_regions:
+                    res_arn = sns.SNSNotification(ctx, imagename)._get_topic_arn(topic_name, region)
+                    topic_arns.append(res_arn)
 
         assert topic_arns == expected
+
+
+@pytest.mark.parametrize(
+    "imagename,regions_in_partition,regions_expected",
+    [
+        ("test-image-10", ["us-east-1", "eu-west-1"], {"topic1": ["us-east-1"]}),
+        (
+            "test-image-11",
+            ["us-east-1", "eu-west-1"],
+            {"topic1": ["us-east-1"], "topic2": []},
+        ),
+        ("test-image-12", ["eu-northwest-1", "ap-southeast-1"], {"topic1": ["eu-northwest-1", "ap-southeast-1"]}),
+    ],
+)
+def test_sns_regions(imagename, regions_in_partition, regions_expected):
+    """
+    Test the regions for a given image
+    """
+    with patch("boto3.client") as bclient_mock:
+        instance = bclient_mock.return_value
+        instance.describe_regions.return_value = {"Regions": [{"RegionName": r} for r in regions_in_partition]}
+        ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+        sns_conf = ctx.conf["images"][imagename]["sns"]
+        instance.list_buckets.return_value = {"Buckets": [{"Name": "bucket1"}]}
+
+        sns_regions = {}
+        for topic in sns_conf:
+            for topic_name, topic_conf in topic.items():
+                sns_regions[topic_name] = sns.SNSNotification(ctx, imagename)._sns_regions(topic_conf)
+
+        assert sns_regions == regions_expected

docs/config-samples/config-minimal-sns.yaml

@@ -17,3 +17,5 @@ awspub:
             subject: "my-topic2-subject"
             message:
               default: "This is message for email protocols. New image $serial is available"
+            regions:
+              - us-east-1

docs/how_to/publish.rst

@@ -242,7 +242,8 @@ SNS Notification
 ~~~~~~~~~~~~~~~~
 
 It's possible to publish messages through the `Simple Notification Service (SNS) <https://docs.aws.amazon.com/sns/latest/dg/welcome.html>`_.
-Delivery to multiple topics is possible, but the topics need to exist in each of the regions an image gets published.
+Delivery to multiple topics is possible, but the topics need to exist in each of the regions where the notification will be sent.
+
 To notify image information to users, the ``sns`` configuration for each image must be filled:
 
 .. literalinclude:: ../config-samples/config-minimal-sns.yaml
@@ -257,6 +258,8 @@ Currently, the supported protocols are ``default`` and ``email`` only, and the `
 send notifications.
 The ``default`` message will be used as a fallback message for any protocols.
 
+Also, Regions can also be specified in ``sns`` configuration to indicate where the notification should be sent. If no regions are specified, SNS will default to using all regions in the partition.
+
 Create the image and use the `publish` command to publish the image and also notify the published images to users:
 
 .. code-block:: shell