Merge pull request #74 from JessicaJang/sns-region-setting

fix: Allow SNS config has separate list of regions

Author: toabctl

awspub/common.py

@@ -1,4 +1,10 @@
-from typing import Tuple
+import logging
+from typing import List, Tuple
+
+import boto3
+from mypy_boto3_ec2.client import EC2Client
+
+logger = logging.getLogger(__name__)
 
 
 def _split_partition(val: str) -> Tuple[str, str]:
@@ -16,3 +22,43 @@ def _split_partition(val: str) -> Tuple[str, str]:
         partition = "aws"
         resource = val
     return partition, resource
+
+
+def _get_regions(region_to_query: str, regions_allowlist: List[str]) -> List[str]:
+    """
+    Get a list of region names querying the `region_to_query` for all regions and
+    then filtering by `regions_allowlist`.
+    If no `regions_allowlist` is given, all queried regions are returned for the
+    current partition.
+    If `regions_allowlist` is given, all regions from that list are returned if
+    the listed region exist in the current partition.
+    Eg. `us-east-1` listed in `regions_allowlist` won't be returned if the current
+    partition is `aws-cn`.
+    :param region_to_query: region name of current partition
+    :type region_to_query: str
+    :praram regions_allowlist: list of regions in config file
+    :type regions_allowlist: List[str]
+    :return: list of regions names
+    :rtype: List[str]
+    """
+
+    # get all available regions
+    ec2client: EC2Client = boto3.client("ec2", region_name=region_to_query)
+    resp = ec2client.describe_regions()
+    ec2_regions_all = [r["RegionName"] for r in resp["Regions"]]
+
+    if regions_allowlist:
+        # filter out regions that are not available in the current partition
+        regions_allowlist_set = set(regions_allowlist)
+        ec2_regions_all_set = set(ec2_regions_all)
+        regions = list(regions_allowlist_set.intersection(ec2_regions_all_set))
+        diff = regions_allowlist_set.difference(ec2_regions_all_set)
+        if diff:
+            logger.warning(
+                f"regions {diff} listed in regions allowlist are not available in the current partition."
+                " Ignoring those."
+            )
+    else:
+        regions = ec2_regions_all
+
+    return regions

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

@@ -9,7 +9,7 @@
 from mypy_boto3_ssm import SSMClient
 
 from awspub import exceptions
-from awspub.common import _split_partition
+from awspub.common import _get_regions, _split_partition
 from awspub.context import Context
 from awspub.image_marketplace import ImageMarketplace
 from awspub.s3 import S3
@@ -121,28 +121,11 @@ def snapshot_name(self) -> str:
     @property
     def image_regions(self) -> List[str]:
         """
-        Get the image regions. Either configured in the image configuration
-        or all available regions.
-        If a region is listed that is not available in the currently used partition,
-        that region will be ignored (eg. having us-east-1 configured but running in the aws-cn
-        partition doesn't include us-east-1 here).
+        Get the image regions.
         """
         if not self._image_regions_cached:
-            # get all available regions
-            ec2client: EC2Client = boto3.client("ec2", region_name=self._s3.bucket_region)
-            resp = ec2client.describe_regions()
-            image_regions_all = [r["RegionName"] for r in resp["Regions"]]
-
-            if self.conf["regions"]:
-                # filter out regions that are not available in the current partition
-                image_regions_configured_set = set(self.conf["regions"])
-                image_regions_all_set = set(image_regions_all)
-                self._image_regions = list(image_regions_configured_set.intersection(image_regions_all_set))
-                diff = image_regions_configured_set.difference(image_regions_all_set)
-                if diff:
-                    logger.warning(f"configured regions {diff} not available in the current partition. Ignoring those.")
-            else:
-                self._image_regions = image_regions_all
+            regions_configured = self.conf["regions"] if "regions" in self.conf else []
+            self._image_regions = _get_regions(self._s3.bucket_region, regions_configured)
             self._image_regions_cached = True
         return self._image_regions
 
@@ -358,14 +341,8 @@ def _sns_publish(self) -> None:
         """
         Publish SNS notifiations about newly available images to subscribers
         """
-        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, region).publish()
+        SNSNotification(self._ctx, self.image_name).publish()
 
     def cleanup(self) -> None:
         """

awspub/sns.py

@@ -11,8 +11,10 @@
 from mypy_boto3_sns.client import SNSClient
 from mypy_boto3_sts.client import STSClient
 
+from awspub.common import _get_regions
 from awspub.context import Context
 from awspub.exceptions import AWSAuthorizationException, AWSNotificationException
+from awspub.s3 import S3
 
 logger = logging.getLogger(__name__)
 
@@ -23,13 +25,13 @@ class SNSNotification(object):
     structuring rules for SNS notification JSON
     """
 
-    def __init__(self, context: Context, image_name: str, region_name: str):
+    def __init__(self, context: Context, image_name: str):
         """
         Construct a message and verify that it is valid
         """
         self._ctx: Context = context
         self._image_name: str = image_name
-        self._region_name: str = region_name
+        self._s3: S3 = S3(context)
 
     @property
     def conf(self) -> List[Dict[str, Any]]:
@@ -38,7 +40,21 @@ 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
+        or all available regions.
+        If a region is listed that is not available in the currently used partition,
+        that region will be ignored (eg. having us-east-1 configured but running in the aws-cn
+        partition doesn't include us-east-1 here).
+        """
+
+        regions_configured = topic_config["regions"] if "regions" in topic_config else []
+        sns_regions = _get_regions(self._s3.bucket_region, regions_configured)
+
+        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
@@ -49,40 +65,41 @@ def _get_topic_arn(self, topic_name: str) -> 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_common.py

@@ -1,6 +1,8 @@
+from unittest.mock import patch
+
 import pytest
 
-from awspub.common import _split_partition
+from awspub.common import _get_regions, _split_partition
 
 
 @pytest.mark.parametrize(
@@ -14,3 +16,19 @@
 )
 def test_common__split_partition(input, expected_output):
     assert _split_partition(input) == expected_output
+
+
+@pytest.mark.parametrize(
+    "regions_in_partition,configured_regions,expected_output",
+    [
+        (["region-1", "region-2"], ["region-1", "region-3"], ["region-1"]),
+        (["region-1", "region-2", "region-3"], ["region-4", "region-5"], []),
+        (["region-1", "region-2"], [], ["region-1", "region-2"]),
+    ],
+)
+def test_common__get_regions(regions_in_partition, configured_regions, expected_output):
+    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]}
+
+        assert _get_regions("", configured_regions) == expected_output

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()