awspub 0.0.1__py3-none-any.whl

awspub/image.py

@@ -0,0 +1,593 @@
+import hashlib
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import boto3
+from mypy_boto3_ec2.client import EC2Client
+from mypy_boto3_ssm import SSMClient
+
+from awspub import exceptions
+from awspub.context import Context
+from awspub.image_marketplace import ImageMarketplace
+from awspub.s3 import S3
+from awspub.snapshot import Snapshot
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _ImageInfo:
+    """
+    Information about a image from EC2
+    """
+
+    image_id: str
+    snapshot_id: Optional[str]
+
+
+class ImageVerificationErrors(str, Enum):
+    """
+    Possible errors for image verification
+    """
+
+    NOT_EXIST = "image does not exist"
+    STATE_NOT_AVAILABLE = "image not available"
+    ROOT_DEVICE_TYPE = "root device type mismatch"
+    ROOT_DEVICE_VOLUME_TYPE = "root device volume type mismatch"
+    ROOT_DEVICE_VOLUME_SIZE = "root device volume size mismatch"
+    ROOT_DEVICE_SNAPSHOT_NOT_COMPLETE = "root device snapshot not complete"
+    BOOT_MODE = "boot mode mismatch"
+    TAGS = "tags mismatch"
+    TPM_SUPPORT = "tpm support mismatch"
+    IMDS_SUPPORT = "imds support mismatch"
+    BILLING_PRODUCTS = "billing products mismatch"
+
+
+class Image:
+    """
+    Handle EC2 Image/AMI API interaction
+    """
+
+    def __init__(self, context: Context, image_name: str):
+        self._ctx: Context = context
+        self._image_name: str = image_name
+        self._image_regions: List[str] = []
+
+        if self._image_name not in self._ctx.conf["images"].keys():
+            raise ValueError(f"image '{self._image_name}' not found in context configuration")
+
+        self._snapshot: Snapshot = Snapshot(context)
+        self._s3: S3 = S3(context)
+
+    def __repr__(self):
+        return f"<{self.__class__} :'{self.image_name}' (snapshot name: {self.snapshot_name})"
+
+    @property
+    def conf(self) -> Dict[str, Any]:
+        """
+        The configuration for the current image (based on "image_name") from context
+        """
+        return self._ctx.conf["images"][self._image_name]
+
+    @property
+    def image_name(self) -> str:
+        """
+        Get the image name
+        """
+        return self._image_name
+
+    @property
+    def snapshot_name(self) -> str:
+        """
+        Get the snapshot name which is a sha256 hexdigest
+
+        The snapshot name is the sha256 hexdigest of the source file given in the source->path
+        configuration option.
+
+        if the "separate_snapshot" config option is set to True, the snapshot name is
+        sha256 hexdigest of the source file given in the source->path conf option and then
+        the sha256 hexdigest of the image-name appended and then the sha256 hexdigest
+        calculated of this concatenated string.
+
+        if the "billing_products" config option is set, the snapshot name is
+        sha256 hexdigest of the source file given in the source->path conf option and then
+        the sha256 hexdigest of each entry in the billing_products appended and then the sha256 hexdigest
+        calculated of this concatenated string.
+
+        Note that both options ("separate_snapshot" and "billing_products") can be combined
+        and the snapshot calculation steps would be combined, too.
+        """
+        s_name = self._ctx.source_sha256
+        if self.conf["separate_snapshot"] is True:
+            s_name += hashlib.sha256(self.image_name.encode("utf-8")).hexdigest()
+
+        if self.conf["billing_products"]:
+            for bp in self.conf["billing_products"]:
+                s_name += hashlib.sha256(bp.encode("utf-8")).hexdigest()
+
+        # in the separate_snapshot and billing_products had no effect, don't do another sha256 of
+        # the source_sha256 to simplify things
+        if s_name == self._ctx.source_sha256:
+            return s_name
+
+        # do a sha256 of the concatenated string
+        return hashlib.sha256(s_name.encode("utf-8")).hexdigest()
+
+    @property
+    def image_regions(self) -> List[str]:
+        """
+        Get the image regions. Either configured in the image configuration
+        or all available regions
+        """
+        if not self._image_regions:
+            if self.conf["regions"]:
+                self._image_regions = self.conf["regions"]
+            else:
+                ec2client: EC2Client = boto3.client("ec2", region_name=self._s3.bucket_region)
+                resp = ec2client.describe_regions()
+                self._image_regions = [r["RegionName"] for r in resp["Regions"]]
+        return self._image_regions
+
+    @property
+    def _tags(self):
+        """
+        Get the tags for this image (common tags + image specific tags)
+        image specific tags override common tags
+        """
+        tags = []
+        # the common tags
+        tags_dict = self._ctx.tags_dict
+        # the image specific tags
+        tags_dict.update(self.conf.get("tags", {}))
+        for name, value in tags_dict.items():
+            tags.append({"Key": name, "Value": value})
+        return tags
+
+    def _share(self, share_conf: List[str], images: Dict[str, _ImageInfo]):
+        """
+        Share images with accounts
+
+        :param share_conf: the share configuration. eg. self.conf["share_create"]
+        :type share_conf: List[str]
+        :param images: a Dict with region names as keys and _ImageInfo objects as values
+        :type images: Dict[str, _ImageInfo]
+        """
+        share_list: List[Dict[str, str]] = [{"UserId": user_id} for user_id in share_conf]
+
+        for region, image_info in images.items():
+            ec2client: EC2Client = boto3.client("ec2", region_name=region)
+            # modify image permissions
+            ec2client.modify_image_attribute(
+                Attribute="LaunchPermission",
+                ImageId=image_info.image_id,
+                LaunchPermission={"Add": share_list},  # type: ignore
+            )
+
+            # modify snapshot permissions
+            if image_info.snapshot_id:
+                ec2client.modify_snapshot_attribute(
+                    Attribute="createVolumePermission",
+                    SnapshotId=image_info.snapshot_id,
+                    CreateVolumePermission={"Add": share_list},  # type: ignore
+                )
+
+        logger.info(f"shared images & snapshots with '{share_conf}'")
+
+    def _get_root_device_snapshot_id(self, image):
+        """
+        Get the root device snapshot id for a given image
+        :param image: a image structure returned by eg. describe_images()["Images"][0]
+        :type image: dict
+        :return: Either None or a snapshot-id
+        :rtype: Optional[str]
+        """
+        root_device_name = image.get("RootDeviceName")
+        if not root_device_name:
+            logger.debug(f"can not get RootDeviceName for image {image}")
+            return None
+        for bdm in image["BlockDeviceMappings"]:
+            if bdm["DeviceName"] == root_device_name:
+                ebs = bdm.get("Ebs")
+                if not ebs:
+                    logger.debug(
+                        f"can not get RootDeviceName. root device {root_device_name} doesn't have a Ebs section"
+                    )
+                    return None
+                logger.debug(f"found Ebs for root device {root_device_name}: {bdm['Ebs']}")
+                return bdm["Ebs"]["SnapshotId"]
+
+    def _get(self, ec2client: EC2Client) -> Optional[_ImageInfo]:
+        """
+        Get the a _ImageInfo for the current image which contains the ami id and
+        root device snapshot id.
+        This relies on the image name to be unique and will raise a MultipleImagesException
+        if multiple images are found.
+
+        :param ec2client: EC2Client
+        :type ec2client: EC2Client
+        :return: Either None or a _ImageInfo
+        :rtype: Optional[_ImageInfo]
+        """
+        resp = ec2client.describe_images(
+            Filters=[
+                {"Name": "name", "Values": [self.image_name]},
+            ],
+            Owners=["self"],
+        )
+
+        if len(resp.get("Images", [])) == 1:
+            root_device_snapshot_id = self._get_root_device_snapshot_id(resp["Images"][0])
+            return _ImageInfo(resp["Images"][0]["ImageId"], root_device_snapshot_id)
+        elif len(resp.get("Images", [])) == 0:
+            return None
+        else:
+            images = [i["ImageId"] for i in resp.get("Images", [])]
+            raise exceptions.MultipleImagesException(
+                f"Found {len(images)} images ({', '.join(images)}) with "
+                f"name {self.image_name} in region {ec2client.meta.region_name}. There should be only 1."
+            )
+
+    def _put_ssm_parameters(self) -> None:
+        """
+        Push the configured SSM parameters to the parameter store
+        """
+        logger.info(f"Pushing SSM parameters for image {self.image_name} in {len(self.image_regions)} regions ...")
+        for region in self.image_regions:
+            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+
+            # image in region not found
+            if not image_info:
+                logger.error(f"image {self.image_name} not available in region {region}. can not push SSM parameter")
+                continue
+
+            ssmclient_region: SSMClient = boto3.client("ssm", region_name=region)
+            # iterate over all defined parameters
+            for parameter in self.conf["ssm_parameter"]:
+                # if overwrite is not allowed, check if the parameter is already there and if so, do nothing
+                if not parameter["allow_overwrite"]:
+                    resp = ssmclient_region.get_parameters(Names=[parameter["name"]])
+                    if len(resp["Parameters"]) >= 1:
+                        # sanity check if the available parameter matches the value we would (but don't) push
+                        if resp["Parameters"][0]["Value"] != image_info.image_id:
+                            logger.warning(
+                                f"SSM parameter {parameter['name']} exists but value does not match "
+                                f"(found {resp['Parameters'][0]['Value']}; expected: {image_info.image_id}"
+                            )
+                        # parameter exists already and overwrite is not allowed so continue
+                        continue
+                # push parameter to store
+                ssmclient_region.put_parameter(
+                    Name=parameter["name"],
+                    Description=parameter.get("description", ""),
+                    Value=image_info.image_id,
+                    Type="String",
+                    Overwrite=parameter["allow_overwrite"],
+                    DataType="aws:ec2:image",
+                    # TODO: tags can't be used together with overwrite
+                    # Tags=self._ctx.tags,
+                )
+
+                logger.info(
+                    f"pushed SSM parameter {parameter['name']} with value {image_info.image_id} in region {region}"
+                )
+
+    def _public(self) -> None:
+        """
+        Make image and underlying root device snapshot public
+        """
+        logger.info(f"Make image {self.image_name} in {len(self.image_regions)} regions public ...")
+
+        for region in self.image_regions:
+            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+            if image_info:
+                ec2client_region.modify_image_attribute(
+                    ImageId=image_info.image_id,
+                    LaunchPermission={
+                        "Add": [
+                            {
+                                "Group": "all",
+                            },
+                        ],
+                    },
+                )
+                logger.info(f"image {image_info.image_id} in region {region} public now")
+
+                if image_info.snapshot_id:
+                    ec2client_region.modify_snapshot_attribute(
+                        SnapshotId=image_info.snapshot_id,
+                        Attribute="createVolumePermission",
+                        GroupNames=[
+                            "all",
+                        ],
+                        OperationType="add",
+                    )
+                    logger.info(
+                        f"snapshot {image_info.snapshot_id} ({image_info.image_id}) in region {region} public now"
+                    )
+                else:
+                    logger.error(
+                        f"snapshot for image {self.image_name} ({image_info.image_id}) not available "
+                        f"in region {region}. can not make public"
+                    )
+            else:
+                logger.error(f"image {self.image_name} not available in region {region}. can not make public")
+
+    def cleanup(self) -> None:
+        """
+        Cleanup/delete the temporary images
+
+        If an image is marked as "temporary" in the configuration, do
+        delete that image in all regions.
+        Note: if a temporary image is public, it won't be deleted. A temporary
+        image should never be public
+        Note: the underlying snapshot is currently not deleted. That might change in
+        the future
+        """
+        if not self.conf["temporary"]:
+            logger.info(f"image {self.image_name} not marked as temporary. no cleanup")
+            return
+
+        # do the cleanup - the image is marked as temporary
+        logger.info(f"Cleanup image {self.image_name} ...")
+        for region in self.image_regions:
+            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+
+            if image_info:
+                resp = ec2client_region.describe_images(
+                    Filters=[
+                        {"Name": "image-id", "Values": [image_info.image_id]},
+                    ]
+                )
+                if resp["Images"][0]["Public"] is True:
+                    # this shouldn't happen because the image is marked as temporary in the config
+                    # so how can it be public?
+                    logger.error(
+                        f"no cleanup for {self.image_name} in {region} because ({image_info.image_id}) image is public"
+                    )
+                else:
+                    ec2client_region.deregister_image(ImageId=image_info.image_id)
+                    logger.info(f"{self.image_name} in {region} ({image_info.image_id}) deleted")
+
+    def list(self) -> Dict[str, _ImageInfo]:
+        """
+        Get image based on the available configuration
+        This doesn't change anything - it just tries to get the available image
+        for the different configured regions
+        :return: a Dict with region names as keys and _ImageInfo objects as values
+        :rtype: Dict[str, _ImageInfo]
+        """
+        images: Dict[str, _ImageInfo] = dict()
+        for region in self.image_regions:
+            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+            if image_info:
+                images[region] = image_info
+            else:
+                logger.warning(f"image {self.image_name} not available in region {region}")
+        return images
+
+    def create(self) -> Dict[str, _ImageInfo]:
+        """
+        Get or create a image based on the available configuration
+
+        :return: a Dict with region names as keys and _ImageInfo objects as values
+        :rtype: Dict[str, _ImageInfo]
+        """
+        # this **must** be the region that is used for S3
+        ec2client: EC2Client = boto3.client("ec2", region_name=self._s3.bucket_region)
+
+        # make sure the initial snapshot exists
+        self._snapshot.create(ec2client, self.snapshot_name)
+
+        # make sure the snapshot exist in all required regions
+        snapshot_ids: Dict[str, str] = self._snapshot.copy(
+            self.snapshot_name, self._s3.bucket_region, self.image_regions
+        )
+
+        images: Dict[str, _ImageInfo] = dict()
+        for region in self.image_regions:
+            ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+            image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+            if image_info:
+                if image_info.snapshot_id != snapshot_ids[region]:
+                    logger.warning(
+                        f"image with name '{self.image_name}' already exists ({image_info.image_id}) "
+                        f"in region {ec2client_region.meta.region_name} but the root device "
+                        f"snapshot id is unexpected (got {image_info.snapshot_id} but expected {snapshot_ids[region]})"
+                    )
+                else:
+                    logger.info(
+                        f"image with name '{self.image_name}' already exists ({image_info.image_id}) "
+                        f"in region {ec2client_region.meta.region_name}"
+                    )
+                images[region] = image_info
+            else:
+                logger.info(
+                    f"creating image with name '{self.image_name}' in "
+                    f"region {ec2client_region.meta.region_name} ..."
+                )
+
+                register_image_kwargs = dict(
+                    Name=self.image_name,
+                    Description=self.conf.get("description", ""),
+                    Architecture=self._ctx.conf["source"]["architecture"],
+                    RootDeviceName=self.conf["root_device_name"],
+                    BlockDeviceMappings=[
+                        {
+                            "Ebs": {
+                                "SnapshotId": snapshot_ids[region],
+                                "VolumeType": self.conf["root_device_volume_type"],
+                                "VolumeSize": self.conf["root_device_volume_size"],
+                            },
+                            "DeviceName": self.conf["root_device_name"],
+                        },
+                        # TODO: make those ephemeral block device mappings configurable
+                        {"VirtualName": "ephemeral0", "DeviceName": "/dev/sdb"},
+                        {"VirtualName": "ephemeral1", "DeviceName": "/dev/sdc"},
+                    ],
+                    EnaSupport=True,
+                    SriovNetSupport="simple",
+                    VirtualizationType="hvm",
+                    BootMode=self.conf["boot_mode"],
+                )
+
+                if self.conf["tpm_support"]:
+                    register_image_kwargs["TpmSupport"] = self.conf["tpm_support"]
+
+                if self.conf["imds_support"]:
+                    register_image_kwargs["ImdsSupport"] = self.conf["imds_support"]
+
+                if self.conf["uefi_data"]:
+                    with open(self.conf["uefi_data"], "r") as f:
+                        uefi_data = f.read()
+                    register_image_kwargs["UefiData"] = uefi_data
+
+                if self.conf["billing_products"]:
+                    register_image_kwargs["BillingProducts"] = self.conf["billing_products"]
+
+                resp = ec2client_region.register_image(**register_image_kwargs)
+                ec2client_region.create_tags(Resources=[resp["ImageId"]], Tags=self._tags)
+                images[region] = _ImageInfo(resp["ImageId"], snapshot_ids[region])
+
+        # wait for the images
+        logger.info(f"Waiting for {len(images)} images to be ready the regions ...")
+        for region, image_info in images.items():
+            ec2client_region_wait: EC2Client = boto3.client("ec2", region_name=region)
+            logger.info(
+                f"Waiting for {image_info.image_id} in {ec2client_region_wait.meta.region_name} "
+                "to exist/be available ..."
+            )
+            waiter_exists = ec2client_region_wait.get_waiter("image_exists")
+            waiter_exists.wait(ImageIds=[image_info.image_id])
+            waiter_available = ec2client_region_wait.get_waiter("image_available")
+            waiter_available.wait(ImageIds=[image_info.image_id])
+        logger.info(f"{len(images)} images are ready")
+
+        # share
+        if self.conf["share"]:
+            self._share(self.conf["share"], images)
+
+        return images
+
+    def publish(self) -> None:
+        """
+        Handle all publication steps
+        - make image and underlying root device snapshot public if the public flag is set
+        - request a new marketplace version for the image in us-east-1 if the marketplace config is present
+        Note: if the temporary flag is set in the image, this method will do nothing
+        Note: this command doesn't unpublish anything!
+        """
+        # never publish temporary images
+        if self.conf["temporary"]:
+            logger.warning(f"image {self.image_name} marked as temporary. do not publish")
+            return
+
+        # make snapshot and image public if requested in the image
+        if self.conf["public"]:
+            self._public()
+        else:
+            logger.info(f"image {self.image_name} not marked as public. do not publish")
+
+        # handle marketplace publication
+        if self.conf["marketplace"]:
+            # the "marketplace" configuration is only valid in the "aws" partition
+            partition = boto3.client("ec2").meta.partition
+            if partition == "aws":
+                logger.info(f"marketplace version request for {self.image_name}")
+                # image needs to be in us-east-1
+                ec2client: EC2Client = boto3.client("ec2", region_name="us-east-1")
+                image_info: Optional[_ImageInfo] = self._get(ec2client)
+                if image_info:
+                    im = ImageMarketplace(self._ctx, self.image_name)
+                    im.request_new_version(image_info.image_id)
+                else:
+                    logger.error(
+                        f"can not request marketplace version for {self.image_name} because no image found in us-east-1"
+                    )
+            else:
+                logger.info(
+                    f"found marketplace config for {self.image_name} and partition 'aws' but "
+                    f"currently using partition {partition}. Ignoring marketplace config."
+                )
+
+        # handle SSM parameter store
+        if self.conf["ssm_parameter"]:
+            self._put_ssm_parameters()
+
+    def _verify(self, region: str) -> List[ImageVerificationErrors]:
+        """
+        Verify (but don't modify or create anything) the image in a single region
+        """
+        problems: List[ImageVerificationErrors] = []
+        ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
+        image_info: Optional[_ImageInfo] = self._get(ec2client_region)
+
+        if not image_info:
+            problems.append(ImageVerificationErrors.NOT_EXIST)
+            return problems
+
+        image_aws = ec2client_region.describe_images(ImageIds=[image_info.image_id])["Images"][0]
+
+        # verify state
+        if image_aws["State"] != "available":
+            problems.append(ImageVerificationErrors.STATE_NOT_AVAILABLE)
+
+        # verify RootDeviceType
+        if image_aws["RootDeviceType"] != "ebs":
+            problems.append(ImageVerificationErrors.ROOT_DEVICE_TYPE)
+
+        # verify BootMode
+        if image_aws["BootMode"] != self.conf["boot_mode"]:
+            problems.append(ImageVerificationErrors.BOOT_MODE)
+
+        # verify RootDeviceVolumeType, RootDeviceVolumeSize and Snapshot
+        for bdm in image_aws["BlockDeviceMappings"]:
+            if bdm.get("DeviceName") and bdm["DeviceName"] == image_aws["RootDeviceName"]:
+                # here's the root device
+                if bdm["Ebs"]["VolumeType"] != self.conf["root_device_volume_type"]:
+                    problems.append(ImageVerificationErrors.ROOT_DEVICE_VOLUME_TYPE)
+                if bdm["Ebs"]["VolumeSize"] != self.conf["root_device_volume_size"]:
+                    problems.append(ImageVerificationErrors.ROOT_DEVICE_VOLUME_SIZE)
+
+                # verify snapshot
+                snapshot_aws = ec2client_region.describe_snapshots(SnapshotIds=[bdm["Ebs"]["SnapshotId"]])["Snapshots"][
+                    0
+                ]
+                if snapshot_aws["State"] != "completed":
+                    problems.append(ImageVerificationErrors.ROOT_DEVICE_SNAPSHOT_NOT_COMPLETE)
+
+        # verify tpm support
+        if self.conf["tpm_support"] and image_aws.get("TpmSupport") != self.conf["tpm_support"]:
+            problems.append(ImageVerificationErrors.TPM_SUPPORT)
+
+        # verify imds support
+        if self.conf["imds_support"] and image_aws.get("ImdsSupport") != self.conf["imds_support"]:
+            problems.append(ImageVerificationErrors.IMDS_SUPPORT)
+
+        # billing products
+        if self.conf["billing_products"] and image_aws.get("BillingProducts") != self.conf["billing_products"]:
+            problems.append(ImageVerificationErrors.BILLING_PRODUCTS)
+
+        # verify tags
+        for tag in image_aws["Tags"]:
+            if tag["Key"] == "Name" and tag["Value"] != self.snapshot_name:
+                problems.append(ImageVerificationErrors.TAGS)
+
+        return problems
+
+    def verify(self) -> Dict[str, List[ImageVerificationErrors]]:
+        """
+        Verify (but don't modify or create anything) that the image configuration
+        matches what is on AWS
+        """
+        logger.info(f"Verifying image {self.image_name} ...")
+        problems: Dict[str, List[ImageVerificationErrors]] = dict()
+        for region in self.image_regions:
+            problems[region] = self._verify(region)
+
+        return problems