awspub 0.0.10__py3-none-any.whl

awspub/sns.py

@@ -0,0 +1,105 @@
+"""
+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.common import _get_regions
+from awspub.context import Context
+from awspub.exceptions import AWSAuthorizationException, AWSNotificationException
+from awspub.s3 import S3
+
+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):
+        """
+        Construct a message and verify that it is valid
+        """
+        self._ctx: Context = context
+        self._image_name: str = image_name
+        self._s3: S3 = S3(context)
+
+    @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 _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
+        :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=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:{region_name}:{account}:{topic_name}"
+
+    def publish(self) -> None:
+        """
+        send notification to subscribers
+        """
+
+        for topic in self.conf:
+            for topic_name, topic_config in topic.items():
+                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",
+                        )
+                    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/config-invalid-s3-extra.yaml

@@ -0,0 +1,12 @@
+awspub:
+  s3:
+    bucket_name: "bucket1"
+    invalid_field: "not allowed" # This is an invalid field
+  source:
+    path: "config1.vmdk"
+    architecture: "x86_64"
+  images:
+    test-image:
+      description: "Test Image"
+      separate_snapshot: "False"
+      boot_mode: "uefi-preferred"

awspub/tests/fixtures/config-minimal.yaml

@@ -0,0 +1,12 @@
+---
+awspub:
+  source:
+    path: "config1.vmdk"
+    architecture: "x86_64"
+
+  s3:
+    bucket_name: "bucket1"
+
+  images:
+    "my-custom-image":
+      boot_mode: "uefi-preferred"

awspub/tests/fixtures/config-valid-nonawspub.yaml

@@ -0,0 +1,13 @@
+awspub:
+  s3:
+    bucket_name: "bucket1"
+  source:
+    path: "config1.vmdk"
+    architecture: "x86_64"
+  images:
+    test-image:
+      description: "Test Image"
+      separate_snapshot: "False"
+      boot_mode: "uefi-preferred"
+notawspub: # to make sure config outside of toplevel `awspub` dict is allowed
+  foo_bar: "irrelevant"

awspub/tests/fixtures/config1.yaml

@@ -0,0 +1,171 @@
+awspub:
+  s3:
+    bucket_name: "bucket1"
+
+  source:
+    # config1.vmdk generated with
+    # dd if=/dev/zero of=config1.raw bs=1K count=1
+    # qemu-img convert -f raw -O vmdk -o subformat=streamOptimized config1.raw config1.vmdk
+    path: "config1.vmdk"
+    architecture: "x86_64"
+
+  images:
+    "test-image-1":
+      description: |
+        A test image
+      boot_mode: "uefi"
+      regions:
+        - region1
+        - region2
+      temporary: true
+      groups:
+        - group1
+        - group2
+    "test-image-2":
+      description: |
+        A test image with a separate snapshot
+      boot_mode: "uefi"
+      separate_snapshot: true
+      groups:
+        - group1
+    "test-image-3":
+      description: |
+        A test image with a separate snapshot and a billing code
+      boot_mode: "uefi"
+      separate_snapshot: true
+      billing_products:
+        - billingcode
+    "test-image-4":
+      description: |
+        A test image without a separate snapshot but a billing product
+      boot_mode: "uefi-preferred"
+      billing_products:
+        - billingcode
+    "test-image-5":
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      boot_mode: "uefi-preferred"
+      billing_products:
+        - billingcode1
+        - billingcode2
+    "test-image-6":
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      boot_mode: "uefi-preferred"
+      regions:
+        - "eu-central-1"
+      public: true
+      tags:
+        key1: value1
+    "test-image-7":
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      boot_mode: "uefi-preferred"
+      regions:
+        - "eu-central-1"
+      public: true
+      temporary: true
+      tags:
+        key2: name
+        name: "not-foobar"
+    "test-image-8":
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      boot_mode: "uefi-preferred"
+      regions:
+        - "eu-central-1"
+        - "us-east-1"
+      public: true
+      tags:
+        key1: value1
+      share:
+        - "123456789123"
+        - "221020170000"
+        - "aws:290620200000"
+        - "aws-cn:334455667788"
+      marketplace:
+        entity_id: "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
+        access_role_arn: "arn:aws:iam::xxxxxxxxxxxx:role/AWSMarketplaceAccess"
+        version_title: "1.0.0"
+        release_notes: "N/A"
+        user_name: "ubuntu"
+        scanning_port: 22
+        os_name: "UBUNTU"
+        os_version: "22.04"
+        usage_instructions: |
+          You can use me
+        recommended_instance_type: "m5.large"
+        security_groups:
+          -
+            from_port: 22
+            ip_protocol: "tcp"
+            ip_ranges:
+              - "0.0.0.0/0"
+            to_port: 22
+      ssm_parameter:
+        -
+          name: /test/image
+        -
+          name: /test/another-image
+    "test-image-9":
+      boot_mode: "uefi"
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      regions:
+        - "eu-central-1"
+        - "us-east-1"
+      ssm_parameter:
+        -
+          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"
+            regions:
+              - "us-east-1"
+    "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"
+            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/fixtures/config2-mapping.yaml

@@ -0,0 +1,2 @@
+key1: "value1"
+key2: "value$2"

awspub/tests/fixtures/config2.yaml

@@ -0,0 +1,48 @@
+awspub:
+  s3:
+    bucket_name: "bucket1"
+
+  source:
+    # config1.vmdk generated with
+    # dd if=/dev/zero of=config1.raw bs=1K count=1
+    # qemu-img convert -f raw -O vmdk -o subformat=streamOptimized config1.raw config1.vmdk
+    path: "config1.vmdk"
+    architecture: "x86_64"
+
+  images:
+    "test-image-$key1":
+      description: |
+        A test image
+      boot_mode: "uefi"
+      regions:
+        - region1
+        - region2
+      temporary: true
+    "test-image-$key2":
+      description: |
+        A test image with a separate snapshot
+      boot_mode: "uefi"
+      separate_snapshot: true
+    "test-image-3":
+      description: |
+        A test image with a separate snapshot and a billing code
+      boot_mode: "uefi"
+      separate_snapshot: true
+      billing_products:
+        - billingcode
+    "test-image-4":
+      description: |
+        A test image without a separate snapshot but a billing product
+      boot_mode: "uefi-preferred"
+      billing_products:
+        - billingcode
+    "test-image-5":
+      description: |
+        A test image without a separate snapshot but multiple billing products
+      boot_mode: "uefi-preferred"
+      billing_products:
+        - billingcode1
+        - billingcode2
+
+  tags:
+    name: "foobar"

awspub/tests/fixtures/config3-duplicate-keys.yaml

@@ -0,0 +1,18 @@
+awspub:
+  s3:
+    bucket_name: "bucket1"
+
+  source:
+    path: "config1.vmdk"
+    architecture: "x86_64"
+
+  images:
+    "test-image-1":
+      description: |
+        A test image
+      boot_mode: "uefi"
+    # second image with the same key
+    "test-image-1":
+      description: |
+        A test image
+      boot_mode: "uefi"

awspub/tests/test_api.py

@@ -0,0 +1,89 @@
+import pathlib
+
+import pytest
+
+from awspub import api, context, image
+
+curdir = pathlib.Path(__file__).parent.resolve()
+
+
+@pytest.mark.parametrize(
+    "group,expected_image_names",
+    [
+        # without any group, all images should be processed
+        (
+            None,
+            [
+                "test-image-1",
+                "test-image-2",
+                "test-image-3",
+                "test-image-4",
+                "test-image-5",
+                "test-image-6",
+                "test-image-7",
+                "test-image-8",
+                "test-image-9",
+                "test-image-10",
+                "test-image-11",
+                "test-image-12",
+            ],
+        ),
+        # with a group that no image as, no image should be processed
+        (
+            "group-not-used",
+            [],
+        ),
+        # with a group that an image has
+        (
+            "group2",
+            ["test-image-1"],
+        ),
+        # with a group that multiple images have
+        (
+            "group1",
+            ["test-image-1", "test-image-2"],
+        ),
+    ],
+)
+def test_api__images_filtered(group, expected_image_names):
+    """
+    Test the _images_filtered() function
+    """
+    ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+
+    image_names = [i[0] for i in api._images_filtered(ctx, group)]
+    assert image_names == expected_image_names
+
+
+@pytest.mark.parametrize(
+    "group,expected",
+    [
+        # without any group, all images should be processed
+        (
+            None,
+            (
+                {"test-image-1": {"eu-central-1": "ami-123", "eu-central-2": "ami-456"}},
+                {
+                    "group1": {"test-image-1": {"eu-central-1": "ami-123", "eu-central-2": "ami-456"}},
+                    "group2": {"test-image-1": {"eu-central-1": "ami-123", "eu-central-2": "ami-456"}},
+                },
+            ),
+        ),
+        # with a group that no image as, image should be there but nothing in the group
+        ("group-not-used", ({"test-image-1": {"eu-central-1": "ami-123", "eu-central-2": "ami-456"}}, {})),
+    ],
+)
+def test_api__images_grouped(group, expected):
+    """
+    Test the _images_grouped() function
+    """
+    ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+    images = [
+        (
+            "test-image-1",
+            image.Image(ctx, "test-image-1"),
+            {"eu-central-1": image._ImageInfo("ami-123", None), "eu-central-2": image._ImageInfo("ami-456", None)},
+        )
+    ]
+    grouped = api._images_grouped(images, group)
+    assert grouped == expected

awspub/tests/test_common.py

@@ -0,0 +1,34 @@
+from unittest.mock import patch
+
+import pytest
+
+from awspub.common import _get_regions, _split_partition
+
+
+@pytest.mark.parametrize(
+    "input,expected_output",
+    [
+        ("123456789123", ("aws", "123456789123")),
+        ("aws:123456789123", ("aws", "123456789123")),
+        ("aws-cn:123456789123", ("aws-cn", "123456789123")),
+        ("aws-us-gov:123456789123", ("aws-us-gov", "123456789123")),
+    ],
+)
+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_context.py

@@ -0,0 +1,88 @@
+import glob
+import os
+import pathlib
+
+import pytest
+from pydantic import ValidationError
+from ruamel.yaml.constructor import DuplicateKeyError
+
+from awspub import context
+
+curdir = pathlib.Path(__file__).parent.resolve()
+
+
+def test_context_create():
+    """
+    Create a Context object from a given configuration
+    """
+    ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+    assert ctx.conf["source"]["path"] == curdir / "fixtures/config1.vmdk"
+    assert ctx.source_sha256 == "6252475408b9f9ee64452b611d706a078831a99b123db69d144d878a0488a0a8"
+    assert ctx.conf["source"]["architecture"] == "x86_64"
+    assert ctx.conf["s3"]["bucket_name"] == "bucket1"
+
+
+def test_context_create_minimal():
+    """
+    Create a Context object from a given minimal configuration
+    """
+    ctx = context.Context(curdir / "fixtures/config-minimal.yaml", None)
+    assert ctx.conf["source"]["path"] == curdir / "fixtures/config1.vmdk"
+    assert ctx.source_sha256 == "6252475408b9f9ee64452b611d706a078831a99b123db69d144d878a0488a0a8"
+    assert ctx.conf["source"]["architecture"] == "x86_64"
+    assert ctx.conf["s3"]["bucket_name"] == "bucket1"
+
+
+def test_context_create_with_mapping():
+    """
+    Create a Context object from a given configuration
+    """
+    ctx = context.Context(curdir / "fixtures/config2.yaml", curdir / "fixtures/config2-mapping.yaml")
+    assert ctx.conf["source"]["path"] == curdir / "fixtures/config1.vmdk"
+    assert ctx.source_sha256 == "6252475408b9f9ee64452b611d706a078831a99b123db69d144d878a0488a0a8"
+    assert ctx.conf["images"].get("test-image-value1")
+    assert ctx.conf["images"].get("test-image-value$2")
+
+
+def test_context_with_docs_config_samples():
+    """
+    Create a Context object with the sample config files used for documentation
+    """
+    config_samples_dir = curdir.parents[1] / "docs" / "config-samples"
+    for f in glob.glob(f"{config_samples_dir}/*.yaml"):
+        mapping_file = f + ".mapping"
+        if os.path.exists(mapping_file):
+            mapping = mapping_file
+        else:
+            mapping = None
+        context.Context(os.path.join(config_samples_dir, f), mapping)
+
+
+def test_context_with_duplicate_image_name():
+    """
+    Create a context with a configuration file that contains a duplicate image name key
+    """
+    with pytest.raises(DuplicateKeyError):
+        context.Context(curdir / "fixtures/config3-duplicate-keys.yaml", None)
+
+
+@pytest.mark.parametrize(
+    "config_file",
+    ["fixtures/config-minimal.yaml", "fixtures/config-valid-nonawspub.yaml"],
+)
+def test_valid_configuration(config_file):
+    """
+    Test with a valid configuration file (no extra fields)
+    """
+    ctx = context.Context(curdir / config_file, None)
+    assert ctx.conf is not None
+    assert ctx.conf["s3"]["bucket_name"] == "bucket1"
+    assert ctx.conf["source"]["architecture"] == "x86_64"
+
+
+def test_invalid_configuration_extra_field():
+    """
+    Test with an invalid configuration file that includes an extra field
+    """
+    with pytest.raises(ValidationError):
+        context.Context(curdir / "fixtures/config-invalid-s3-extra.yaml", None)