Merge pull request #71 from JessicaJang/sns-email-notification

Add Notification API call

Author: JessicaJang

.custom_wordlist.txt

@@ -14,3 +14,4 @@ https
 io
 readthedocs
 vmdk
+SNS

awspub/configmodels.py

@@ -1,4 +1,5 @@
 import pathlib
+from enum import Enum
 from typing import Dict, List, Literal, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -94,6 +95,34 @@ class ConfigImageSSMParameterModel(BaseModel):
     )
 
 
+class SNSNotificationProtocol(str, Enum):
+    DEFAULT = "default"
+    EMAIL = "email"
+
+
+class ConfigImageSNSNotificationModel(BaseModel):
+    """
+    Image/AMI SNS Notification specific configuration to notify subscribers about new images availability
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    subject: str = Field(description="The subject of SNS Notification", min_length=1, max_length=99)
+    message: Dict[SNSNotificationProtocol, str] = Field(
+        description="The body of the message to be sent to subscribers.",
+        default={SNSNotificationProtocol.DEFAULT: ""},
+    )
+
+    @field_validator("message")
+    def check_message(cls, value):
+        # Check message protocols have default key
+        # Message should contain at least a top-level JSON key of “default”
+        # with a value that is a string
+        if SNSNotificationProtocol.DEFAULT not in value:
+            raise ValueError(f"{SNSNotificationProtocol.DEFAULT.value} key is required to send SNS notification")
+        return value
+
+
 class ConfigImageModel(BaseModel):
     """
     Image/AMI configuration.
@@ -148,6 +177,9 @@ class ConfigImageModel(BaseModel):
     )
     groups: Optional[List[str]] = Field(description="Optional list of groups this image is part of", default=[])
     tags: Optional[Dict[str, str]] = Field(description="Optional Tags to apply to this image only", default={})
+    sns: Optional[List[Dict[str, ConfigImageSNSNotificationModel]]] = Field(
+        description="Optional list of SNS Notification related configuration", default=None
+    )
 
     @field_validator("share")
     @classmethod

awspub/exceptions.py

@@ -14,3 +14,11 @@ class BucketDoesNotExistException(Exception):
     def __init__(self, bucket_name: str, *args, **kwargs):
         msg = f"The bucket named '{bucket_name}' does not exist. You will need to create the bucket before proceeding."
         super().__init__(msg, *args, **kwargs)
+
+
+class AWSNotificationException(Exception):
+    pass
+
+
+class AWSAuthorizationException(Exception):
+    pass

awspub/image.py

@@ -14,6 +14,7 @@
 from awspub.image_marketplace import ImageMarketplace
 from awspub.s3 import S3
 from awspub.snapshot import Snapshot
+from awspub.sns import SNSNotification
 
 logger = logging.getLogger(__name__)
 
@@ -353,6 +354,19 @@ def _public(self) -> None:
             else:
                 logger.error(f"image {self.image_name} not available in region {region}. can not make public")
 
+    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()
+
     def cleanup(self) -> None:
         """
         Cleanup/delete the temporary images
@@ -556,6 +570,10 @@ def publish(self) -> None:
                     f"currently using partition {partition}. Ignoring marketplace config."
                 )
 
+        # send ssn notification
+        if self.conf["sns"]:
+            self._sns_publish()
+
     def _verify(self, region: str) -> List[ImageVerificationErrors]:
         """
         Verify (but don't modify or create anything) the image in a single region

awspub/sns.py

@@ -0,0 +1,88 @@
+"""
+Methods used to handle notifications for AWS using SNS
+"""
+
+import json
+import logging
+from typing import Any, Dict, List
+
+import boto3
+from botocore.exceptions import ClientError
+from mypy_boto3_sns.client import SNSClient
+from mypy_boto3_sts.client import STSClient
+
+from awspub.context import Context
+from awspub.exceptions import AWSAuthorizationException, AWSNotificationException
+
+logger = logging.getLogger(__name__)
+
+
+class SNSNotification(object):
+    """
+    A data object that contains validation logic and
+    structuring rules for SNS notification JSON
+    """
+
+    def __init__(self, context: Context, image_name: str, region_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
+
+    @property
+    def conf(self) -> List[Dict[str, Any]]:
+        """
+        The sns configuration for the current image (based on "image_name") from context
+        """
+        return self._ctx.conf["images"][self._image_name]["sns"]
+
+    def _get_topic_arn(self, topic_name: str) -> str:
+        """
+        Calculate topic ARN based on partition, region, account and topic name
+        :param topic_name: Name of topic
+        :type topic_name: str
+        :param region_name: name of region
+        :type region_name: str
+        :return: return topic ARN
+        :rtype: str
+        """
+
+        stsclient: STSClient = boto3.client("sts", region_name=self._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}"
+
+    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."
+                        )
+                    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."
+                )

awspub/tests/fixtures/config1.yaml

@@ -118,6 +118,35 @@ awspub:
         -
           name: /awspub-test/param2
           allow_overwrite: true
+    "test-image-10":
+      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"
+    "test-image-11":
+      boot_mode: "uefi"
+      description: |
+        A test image without a separate snapshot but multiple sns configs
+      regions:
+        - "us-east-1"
+        - "eu-central-1"
+      sns:
+        - "topic1":
+            subject: "topic1-subject"
+            message:
+              default: "default-message"
+              email: "email-message"
+        - "topic2":
+            subject: "topic2-subject"
+            message:
+              default: "default-message"
 
   tags:
     name: "foobar"

awspub/tests/test_api.py

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

awspub/tests/test_image.py

@@ -127,16 +127,32 @@ def test_image___get_root_device_snapshot_id(root_device_name, block_device_mapp
 
 
 @pytest.mark.parametrize(
-    "imagename,partition,called_mod_image,called_mod_snapshot,called_start_change_set,called_put_parameter",
+    (
+        "imagename",
+        "partition",
+        "called_mod_image",
+        "called_mod_snapshot",
+        "called_start_change_set",
+        "called_put_parameter",
+        "called_sns_publish",
+    ),
     [
-        ("test-image-6", "aws", True, True, False, False),
-        ("test-image-7", "aws", False, False, False, False),
-        ("test-image-8", "aws", True, True, True, True),
-        ("test-image-8", "aws-cn", True, True, False, True),
+        ("test-image-6", "aws", True, True, False, False, False),
+        ("test-image-7", "aws", False, False, False, False, False),
+        ("test-image-8", "aws", True, True, True, True, False),
+        ("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),
     ],
 )
 def test_image_publish(
-    imagename, partition, called_mod_image, called_mod_snapshot, called_start_change_set, called_put_parameter
+    imagename,
+    partition,
+    called_mod_image,
+    called_mod_snapshot,
+    called_start_change_set,
+    called_put_parameter,
+    called_sns_publish,
 ):
     """
     Test the publish() for a given image
@@ -167,13 +183,15 @@ 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()
         assert instance.modify_image_attribute.called == called_mod_image
         assert instance.modify_snapshot_attribute.called == called_mod_snapshot
         assert instance.start_change_set.called == called_start_change_set
         assert instance.put_parameter.called == called_put_parameter
+        assert instance.publish.called == called_sns_publish
 
 
 def test_image__get_zero_images():