awspub 0.0.10__py3-none-any.whl

awspub/image_marketplace.py

@@ -0,0 +1,120 @@
+import logging
+import re
+from typing import Any, Dict
+
+import boto3
+from mypy_boto3_marketplace_catalog import MarketplaceCatalogClient
+
+from awspub.context import Context
+
+logger = logging.getLogger(__name__)
+
+
+class ImageMarketplace:
+    """
+    Handle AWS Marketplace API interaction
+    """
+
+    def __init__(self, context: Context, image_name: str):
+        self._ctx: Context = context
+        self._image_name: str = image_name
+        # marketplace-catalog API is only available via us-east-1
+        self._mpclient: MarketplaceCatalogClient = boto3.client("marketplace-catalog", region_name="us-east-1")
+
+    @property
+    def conf(self) -> Dict[str, Any]:
+        """
+        The marketplace configuration for the current image (based on "image_name") from context
+        """
+        return self._ctx.conf["images"][self._image_name]["marketplace"]
+
+    def request_new_version(self, image_id: str) -> None:
+        """
+        Request a new Marketplace version for the given image Id
+
+        :param image_id: an image Id (in the format 'ami-123')
+        :type image_id: str
+        """
+        entity = self._mpclient.describe_entity(Catalog="AWSMarketplace", EntityId=self.conf["entity_id"])
+        # check if the version already exists
+        for version in entity["DetailsDocument"]["Versions"]:
+            if version["VersionTitle"] == self.conf["version_title"]:
+                logger.info(f"Marketplace version '{self.conf['version_title']}' already exists. Do nothing")
+                return
+
+        # version doesn't exist already - create a new one
+        changeset = self._request_new_version_changeset(image_id)
+        changeset_name = ImageMarketplace.sanitize_changeset_name(
+            f"New version request for {self.conf['version_title']}"
+        )
+        resp = self._mpclient.start_change_set(
+            Catalog="AWSMarketplace", ChangeSet=changeset, ChangeSetTags=self._ctx.tags, ChangeSetName=changeset_name
+        )
+        logger.info(
+            f"new version '{self.conf['version_title']}' (image: {image_id}) for entity "
+            f"{self.conf['entity_id']} requested (changeset-id: {resp['ChangeSetId']})"
+        )
+
+    def _request_new_version_changeset(self, image_id: str):
+        """
+        Create a changeset structure for a new AmiProduct version
+        See https://docs.aws.amazon.com/marketplace-catalog/latest/api-reference/ami-products.html#ami-add-version
+
+        :param image_id: an image Id (in the format 'ami-123')
+        :type image_id: str
+        :return: A changeset structure to request a new version
+        :rtype: List[Dict[str, Any]]
+        """
+        return [
+            {
+                "ChangeType": "AddDeliveryOptions",
+                "Entity": {
+                    "Identifier": self.conf["entity_id"],
+                    "Type": "AmiProduct@1.0",
+                },
+                "DetailsDocument": {
+                    "Version": {
+                        "VersionTitle": self.conf["version_title"],
+                        "ReleaseNotes": self.conf["release_notes"],
+                    },
+                    "DeliveryOptions": [
+                        {
+                            "Details": {
+                                "AmiDeliveryOptionDetails": {
+                                    "AmiSource": {
+                                        "AmiId": image_id,
+                                        "AccessRoleArn": self.conf["access_role_arn"],
+                                        "UserName": self.conf["user_name"],
+                                        "OperatingSystemName": self.conf["os_name"],
+                                        "OperatingSystemVersion": self.conf["os_version"],
+                                    },
+                                    "UsageInstructions": self.conf["usage_instructions"],
+                                    "RecommendedInstanceType": self.conf["recommended_instance_type"],
+                                    "SecurityGroups": [
+                                        {
+                                            "IpProtocol": sg["ip_protocol"],
+                                            "IpRanges": [ipr for ipr in sg["ip_ranges"]],
+                                            "FromPort": sg["from_port"],
+                                            "ToPort": sg["to_port"],
+                                        }
+                                        for sg in self.conf["security_groups"]
+                                    ],
+                                }
+                            }
+                        }
+                    ],
+                },
+            }
+        ]
+
+    @staticmethod
+    def sanitize_changeset_name(name: str) -> str:
+        # changeset names can only include alphanumeric characters, whitespace, and any combination of the following
+        # characters: _+=.:@- This regex pattern takes the list of allowed characters, does a negative match on the
+        # string and removes all matched (i.e. disallowed) characters. See [0] for reference.
+        # [0] https://docs.aws.amazon.com/marketplace-catalog/latest/api-reference/API_StartChangeSet.html#API_StartChangeSet_RequestSyntax  # noqa
+        return re.sub(
+            "[^\\w\\s+=.:@-]",
+            "",
+            name,
+        )

awspub/s3.py

@@ -0,0 +1,262 @@
+import base64
+import hashlib
+import logging
+import os
+from typing import Dict
+
+import boto3
+from mypy_boto3_s3.type_defs import CompletedPartTypeDef
+
+from awspub.context import Context
+from awspub.exceptions import BucketDoesNotExistException
+
+# chunk size is required for calculating the checksums
+MULTIPART_CHUNK_SIZE = 8 * 1024 * 1024
+
+
+logger = logging.getLogger(__name__)
+
+
+class S3:
+    """
+    Handle S3 API interaction
+    """
+
+    def __init__(self, context: Context):
+        """
+        :param context:
+        "type context: awspub.context.Context
+        """
+        self._ctx: Context = context
+        self._s3client = boto3.client("s3")
+        self._bucket_region = None
+
+    @property
+    def bucket_region(self):
+        if not self._bucket_region:
+            if not self._bucket_exists():
+                raise BucketDoesNotExistException(self.bucket_name)
+            self._bucket_region = self._s3client.head_bucket(Bucket=self.bucket_name)["BucketRegion"]
+
+        return self._bucket_region
+
+    @property
+    def bucket_name(self):
+        return self._ctx.conf["s3"]["bucket_name"]
+
+    def __repr__(self):
+        return (
+            f"<{self.__class__} bucket:'{self.bucket_name}' "
+            f"region:'{self.bucket_region} key:{self._ctx.source_sha256}'>"
+        )
+
+    def _multipart_sha256sum(self, file_path: str) -> str:
+        """
+        Calculate the sha256 checksum like AWS does it (in a multipart upload) per chunk
+        See https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html#large-object-checksums
+
+        :param file_path: the path to the local file to upload
+        :type file_path: str
+        """
+        sha256_list = []
+        count = 0
+        with open(file_path, "rb") as f:
+            for chunk in iter(lambda: f.read(MULTIPART_CHUNK_SIZE), b""):
+                sha256_list.append(hashlib.sha256(chunk))
+                count += 1
+
+        sha256_list_digest_concatenated = b"".join([s.digest() for s in sha256_list])
+        sha256_b64 = base64.b64encode(hashlib.sha256(sha256_list_digest_concatenated).digest())
+        return f"{sha256_b64.decode('ascii')}-{count}"
+
+    def _bucket_exists(self) -> bool:
+        """
+        Check if the S3 bucket from context exists
+
+        :return: True if the bucket exists, otherwise False
+        :rtype: bool
+        """
+        resp = self._s3client.list_buckets()
+        return self.bucket_name in [b["Name"] for b in resp["Buckets"]]
+
+    def upload_file(self, source_path: str):
+        """
+        Upload a given file to the bucket from context. The key name will be the sha256sum hexdigest of the file.
+        If a file with that name already exist in the given bucket and the calculated sha256sum matches
+        the sha256sum from S3, nothing will be uploaded. Instead the existing file will be used.
+        This method does use a multipart upload internally so an upload can be retriggered in case
+        of errors and the previously uploaded content will be reused.
+        Note: be aware that failed multipart uploads are not deleted. So it's recommended to setup
+        a bucket lifecycle rule to delete incomplete multipart uploads.
+        See https://docs.aws.amazon.com/AmazonS3/latest/userguide//mpu-abort-incomplete-mpu-lifecycle-config.html
+
+        :param source_path: the path to the local file to upload (usually a .vmdk file)
+        :type source_path: str
+        """
+        # make sure the bucket exists
+        if not self._bucket_exists():
+            raise BucketDoesNotExistException(self.bucket_name)
+
+        s3_sha256sum = self._multipart_sha256sum(source_path)
+
+        try:
+            # check if the key exists already in the bucket and if so, if the multipart upload
+            # sha256sum does match
+            head = self._s3client.head_object(
+                Bucket=self.bucket_name, Key=self._ctx.source_sha256, ChecksumMode="ENABLED"
+            )
+
+            if head["ChecksumSHA256"] == s3_sha256sum:
+                logger.info(
+                    f"'{self._ctx.source_sha256}' in bucket '{self.bucket_name}' "
+                    "already exists and sha256sum matches. nothing to upload to S3"
+                )
+                return
+            else:
+                logger.warning(
+                    f"'{self._ctx.source_sha256}' in bucket '{self.bucket_name}' "
+                    f"already exists but sha256sum does not match. Will be overwritten ..."
+                )
+        except Exception:
+            logging.debug(f"Can not find '{self._ctx.source_sha256}' in bucket '{self.bucket_name}'")
+
+        # do the real upload
+        self._upload_file_multipart(source_path, s3_sha256sum)
+
+    def _get_multipart_upload_id(self) -> str:
+        """
+        Get an existing or create a multipart upload id
+
+        :return: a multipart upload id
+        :rtype: str
+        """
+        resp = self._s3client.list_multipart_uploads(Bucket=self.bucket_name)
+        multipart_uploads = [
+            upload["UploadId"] for upload in resp.get("Uploads", []) if upload["Key"] == self._ctx.source_sha256
+        ]
+        if len(multipart_uploads) == 1:
+            logger.info(f"found existing multipart upload '{multipart_uploads[0]}' for key '{self._ctx.source_sha256}'")
+            return multipart_uploads[0]
+        elif len(multipart_uploads) == 0:
+            # create a new multipart upload
+            resp_create = self._s3client.create_multipart_upload(
+                Bucket=self.bucket_name,
+                Key=self._ctx.source_sha256,
+                ChecksumAlgorithm="SHA256",
+                ACL="private",
+            )
+            upload_id = resp_create["UploadId"]
+            logger.info(
+                f"new multipart upload (upload id: '{upload_id})' started in bucket "
+                f"{self.bucket_name} for key {self._ctx.source_sha256}"
+            )
+            # if there's an expire rule configured for that bucket, inform about it
+            if resp_create.get("AbortDate"):
+                logger.info(
+                    f"multipart upload '{upload_id}' will expire at "
+                    f"{resp_create['AbortDate']} (rule: {resp_create.get('AbortRuleId')})"
+                )
+            else:
+                logger.warning("there is no matching expire/lifecycle rule configured for incomplete multipart uploads")
+            return upload_id
+        else:
+            # multiple multipart uploads for the same key available
+            logger.warning(
+                f"there are multiple ({len(multipart_uploads)}) multipart uploads ongoing in "
+                f"bucket {self.bucket_name} for key {self._ctx.source_sha256}"
+            )
+            logger.warning("using the first found multipart upload but you should delete pending multipart uploads")
+            return multipart_uploads[0]
+
+    def _upload_file_multipart(self, source_path: str, s3_sha256sum: str) -> None:
+        """
+        Upload a given file to the bucket from context. The key name will be the sha256sum hexdigest of the file
+
+        :param source_path: the path to the local file to upload (usually a .vmdk file)
+        :type source_path: str
+        :param s3_sha256sum: the sha256sum how S3 calculates it
+        :type s3_sha256sum: str
+        """
+        upload_id = self._get_multipart_upload_id()
+
+        logger.info(f"using upload id '{upload_id}' for multipart upload of '{source_path}' ...")
+        resp_list_parts = self._s3client.list_parts(
+            Bucket=self.bucket_name, Key=self._ctx.source_sha256, UploadId=upload_id
+        )
+
+        # sanity check for the used checksum algorithm
+        if resp_list_parts["ChecksumAlgorithm"] != "SHA256":
+            logger.error(f"available ongoing multipart upload '{upload_id}' does not use SHA256 as checksum algorithm")
+
+        # already available parts
+        parts_available = {p["PartNumber"]: p for p in resp_list_parts.get("Parts", [])}
+        # keep a list of parts (either already available or created) required to complete the multipart upload
+        parts: Dict[int, CompletedPartTypeDef] = {}
+        parts_size_done: int = 0
+        source_path_size: int = os.path.getsize(source_path)
+        with open(source_path, "rb") as f:
+            # parts start at 1 (not 0)
+            for part_number, chunk in enumerate(iter(lambda: f.read(MULTIPART_CHUNK_SIZE), b""), start=1):
+                # the sha256sum of the current part
+                sha256_part = base64.b64encode(hashlib.sha256(chunk).digest()).decode("ascii")
+                # do nothing if that part number already exist and the sha256sum matches
+                if parts_available.get(part_number):
+                    if parts_available[part_number]["ChecksumSHA256"] == sha256_part:
+                        logger.info(f"part {part_number} already exists and sha256sum matches. continue")
+                        parts[part_number] = dict(
+                            PartNumber=part_number,
+                            ETag=parts_available[part_number]["ETag"],
+                            ChecksumSHA256=parts_available[part_number]["ChecksumSHA256"],
+                        )
+                        parts_size_done += len(chunk)
+                        continue
+                    else:
+                        logger.info(f"part {part_number} already exists but will be overwritten")
+
+                # upload a new part
+                resp_upload_part = self._s3client.upload_part(
+                    Body=chunk,
+                    Bucket=self.bucket_name,
+                    ContentLength=len(chunk),
+                    ChecksumAlgorithm="SHA256",
+                    ChecksumSHA256=sha256_part,
+                    Key=self._ctx.source_sha256,
+                    PartNumber=part_number,
+                    UploadId=upload_id,
+                )
+                parts_size_done += len(chunk)
+                # add new part to the dict of parts
+                parts[part_number] = dict(
+                    PartNumber=part_number,
+                    ETag=resp_upload_part["ETag"],
+                    ChecksumSHA256=sha256_part,
+                )
+                logger.info(
+                    f"part {part_number} uploaded ({round(parts_size_done/source_path_size * 100, 2)}% "
+                    f"; {parts_size_done} / {source_path_size} bytes)"
+                )
+
+        logger.info(
+            f"finishing the multipart upload for key '{self._ctx.source_sha256}' in bucket {self.bucket_name} now ..."
+        )
+        # finish the multipart upload
+        self._s3client.complete_multipart_upload(
+            Bucket=self.bucket_name,
+            Key=self._ctx.source_sha256,
+            UploadId=upload_id,
+            ChecksumSHA256=s3_sha256sum,
+            MultipartUpload={"Parts": [value for key, value in parts.items()]},
+        )
+        logger.info(
+            f"multipart upload finished and key '{self._ctx.source_sha256}' now "
+            f"available in bucket '{self.bucket_name}'"
+        )
+
+        # add tagging to the final s3 object
+        self._s3client.put_object_tagging(
+            Bucket=self.bucket_name,
+            Key=self._ctx.source_sha256,
+            Tagging={
+                "TagSet": self._ctx.tags,
+            },
+        )

awspub/snapshot.py

@@ -0,0 +1,241 @@
+import logging
+from typing import Dict, List, Optional
+
+import boto3
+from mypy_boto3_ec2.client import EC2Client
+
+from awspub import exceptions
+from awspub.context import Context
+
+logger = logging.getLogger(__name__)
+
+
+class Snapshot:
+    """
+    Handle EC2 Snapshot API interaction
+    """
+
+    def __init__(self, context: Context):
+        self._ctx: Context = context
+
+    def _get(self, ec2client: EC2Client, snapshot_name: str) -> Optional[str]:
+        """
+        Get the snapshot id for the given name or None
+
+        :param ec2client: EC2 client for a specific region
+        :type ec2client: EC2Client
+        :param snapshot_name: the Snapshot name
+        :type snapshot_name: str
+        :return: Either None or a snapshot-id
+        :rtype: Optional[str]
+        """
+        resp = ec2client.describe_snapshots(
+            Filters=[
+                {
+                    "Name": "tag:Name",
+                    "Values": [
+                        snapshot_name,
+                    ],
+                },
+                {
+                    "Name": "status",
+                    "Values": [
+                        "pending",
+                        "completed",
+                    ],
+                },
+            ],
+            OwnerIds=["self"],
+        )
+        if len(resp.get("Snapshots", [])) == 1:
+            return resp["Snapshots"][0]["SnapshotId"]
+        elif len(resp.get("Snapshots", [])) == 0:
+            return None
+        else:
+            raise exceptions.MultipleSnapshotsException(
+                f"Found {len(resp.get('Snapshots', []))} snapshots with "
+                f"name '{snapshot_name}' in region {ec2client.meta.region_name}"
+            )
+
+    def _get_import_snapshot_task(self, ec2client: EC2Client, snapshot_name: str) -> Optional[str]:
+        """
+        Get a import snapshot task for the given name
+
+        :param ec2client: EC2 client for a specific region
+        :type ec2client: EC2Client
+        :param snapshot_name: the Snapshot name
+        :type snapshot_name: str
+        :return: Either None or a import-snapshot-task-id
+        :rtype: Optional[str]
+        """
+        resp = ec2client.describe_import_snapshot_tasks(
+            Filters=[
+                {
+                    "Name": "tag:Name",
+                    "Values": [
+                        snapshot_name,
+                    ],
+                }
+            ]
+        )
+        # API doesn't support filters by status so filter here
+        tasks: List = resp.get("ImportSnapshotTasks", [])
+        # we already know here that the snapshot does not exist (checked in create() before calling this
+        # function). so ignore "deleted" or "completed" tasks here
+        # it might happen (for whatever reason) that a task got completed but the snapshot got deleted
+        # afterwards. In that case a "completed" task for the given snapshot_name exists but
+        # that doesn't help so ignore it
+        tasks = [t for t in tasks if t["SnapshotTaskDetail"]["Status"] not in ["deleted", "completed"]]
+        if len(tasks) == 1:
+            return tasks[0]["ImportTaskId"]
+        elif len(tasks) == 0:
+            return None
+        else:
+            raise exceptions.MultipleImportSnapshotTasksException(
+                f"Found {len(tasks)} import snapshot tasks with "
+                f"name '{snapshot_name}' in region {ec2client.meta.region_name}"
+            )
+
+    def create(self, ec2client: EC2Client, snapshot_name: str) -> str:
+        """
+        Create a EC2 snapshot with the given name
+        If the snapshot already exists, just return the snapshot-id for the existing snapshot.
+
+        :param ec2client: EC2 client for a specific region
+        :type ec2client: EC2Client
+        :param snapshot_name: the Snapshot name
+        :type snapshot_name: str
+        :return: a snapshot-id
+        :rtype: str
+        """
+        # does a snapshot with the given name already exists?
+        snap_id: Optional[str] = self._get(ec2client, snapshot_name)
+        if snap_id:
+            logger.info(f"snapshot with name '{snapshot_name}' already exists in region {ec2client.meta.region_name}")
+            return snap_id
+
+        logger.info(
+            f"Create snapshot from bucket '{self._ctx.conf['s3']['bucket_name']}' "
+            f"for '{snapshot_name}' in region  {ec2client.meta.region_name}"
+        )
+
+        # extend tags
+        tags = self._ctx.tags
+        tags.append({"Key": "Name", "Value": snapshot_name})
+
+        # does a import snapshot task with the given name already exist?
+        import_snapshot_task_id: Optional[str] = self._get_import_snapshot_task(ec2client, snapshot_name)
+        if import_snapshot_task_id:
+            logger.info(
+                f"import snapshot task ({import_snapshot_task_id}) with "
+                f"name '{snapshot_name}' exists in region {ec2client.meta.region_name}"
+            )
+        else:
+            resp = ec2client.import_snapshot(
+                Description="Import ",
+                DiskContainer={
+                    "Description": "",
+                    "Format": "vmdk",
+                    "UserBucket": {
+                        "S3Bucket": self._ctx.conf["s3"]["bucket_name"],
+                        "S3Key": self._ctx.source_sha256,
+                    },
+                },
+                TagSpecifications=[
+                    {"ResourceType": "import-snapshot-task", "Tags": tags},
+                ],
+            )
+            import_snapshot_task_id = resp["ImportTaskId"]
+
+        logger.info(
+            f"Waiting for snapshot import task (id: {import_snapshot_task_id}) "
+            f"in region {ec2client.meta.region_name} ..."
+        )
+
+        waiter_import = ec2client.get_waiter("snapshot_imported")
+        waiter_import.wait(ImportTaskIds=[import_snapshot_task_id], WaiterConfig={"Delay": 30, "MaxAttempts": 90})
+
+        task_details = ec2client.describe_import_snapshot_tasks(ImportTaskIds=[import_snapshot_task_id])
+        snapshot_id = task_details["ImportSnapshotTasks"][0]["SnapshotTaskDetail"]["SnapshotId"]
+
+        # create tags before waiting for completion so the tags are already there
+        ec2client.create_tags(Resources=[snapshot_id], Tags=tags)
+
+        waiter_completed = ec2client.get_waiter("snapshot_completed")
+        waiter_completed.wait(SnapshotIds=[snapshot_id], WaiterConfig={"Delay": 30, "MaxAttempts": 60})
+
+        logger.info(f"Snapshot import as '{snapshot_id}' in region {ec2client.meta.region_name} done")
+        return snapshot_id
+
+    def _copy(self, snapshot_name: str, source_region: str, destination_region: str) -> str:
+        """
+        Copy a EC2 snapshot for the given context to the destination region
+        NOTE: we don't wait for the snapshot to complete here!
+
+        :param snapshot_name: the Snapshot name to copy
+        :type snapshot_name: str
+        :param source_region: a region to copy the snapshot from
+        :type source_region: str
+        :param destination_region: a region to copy the snapshot to
+        :type destionation_region: str
+
+        :return: the existing or created snapshot-id
+        :rtype: str
+        """
+
+        # does the snapshot with that name already exist in the destination region?
+        ec2client_dest: EC2Client = boto3.client("ec2", region_name=destination_region)
+        snapshot_id: Optional[str] = self._get(ec2client_dest, snapshot_name)
+        if snapshot_id:
+            logger.info(
+                f"snapshot with name '{snapshot_name}' already "
+                f"exists ({snapshot_id}) in destination region {ec2client_dest.meta.region_name}"
+            )
+            return snapshot_id
+
+        ec2client_source: EC2Client = boto3.client("ec2", region_name=source_region)
+        source_snapshot_id: Optional[str] = self._get(ec2client_source, snapshot_name)
+        if not source_snapshot_id:
+            raise ValueError(
+                f"Can not find source snapshot with name '{snapshot_name}' "
+                f"in region {ec2client_source.meta.region_name}"
+            )
+
+        logger.info(f"Copy snapshot {source_snapshot_id} from " f"{source_region} to {destination_region}")
+        # extend tags
+        tags = self._ctx.tags
+        tags.append({"Key": "Name", "Value": snapshot_name})
+        resp = ec2client_dest.copy_snapshot(
+            SourceRegion=source_region,
+            SourceSnapshotId=source_snapshot_id,
+            TagSpecifications=[{"ResourceType": "snapshot", "Tags": tags}],
+        )
+
+        # note: we don't wait for the snapshot to complete here!
+        return resp["SnapshotId"]
+
+    def copy(self, snapshot_name: str, source_region: str, destination_regions: List[str]) -> Dict[str, str]:
+        """
+        Copy a snapshot to multiple regions
+
+        :param snapshot_name: the Snapshot name to copy
+        :type snapshot_name: str
+        :param source_region: a region to copy the snapshot from
+        :type source_region: str
+        :param destination_regions: a list of regions to copy the snaphot to
+        :type destionation_regions: List[str]
+        :return: a dict with region/snapshot-id mapping for the newly copied snapshots
+        :rtype: Dict[str, str] where the key is a region name and the value a snapshot-id
+        """
+        snapshot_ids: Dict[str, str] = dict()
+        for destination_region in destination_regions:
+            snapshot_ids[destination_region] = self._copy(snapshot_name, source_region, destination_region)
+
+        logger.info(f"Waiting for {len(snapshot_ids)} snapshots to appear in the destination regions ...")
+        for destination_region, snapshot_id in snapshot_ids.items():
+            ec2client_dest = boto3.client("ec2", region_name=destination_region)
+            waiter = ec2client_dest.get_waiter("snapshot_completed")
+            logger.info(f"Waiting for {snapshot_id} in {ec2client_dest.meta.region_name} to complete ...")
+            waiter.wait(SnapshotIds=[snapshot_id], WaiterConfig={"Delay": 30, "MaxAttempts": 90})
+
+        return snapshot_ids