awspub-register

#!/usr/bin/env python3

"""
Helper using awspub to register an image on a single region.


Copyright (C) <2025> <Thomas Bechtold>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as
published by the Free Software Foundation, version 3.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty
of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>. 
"""

import hashlib
import logging
import argparse
import awspub
import pathlib
import tempfile
import yaml


logger = logging.getLogger(__name__)


def _sha256sum(file_path: pathlib.Path):
    """
    Calculate a sha256 sum for a given file
    """
    sha256_hash = hashlib.sha256()
    with open(file_path.resolve(), "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash


def _config_create(args):
    """
    create a configuration which can be used to register an image
    """
    source_sha256 = _sha256sum(args.source).hexdigest()
    config = {
        "awspub": {
            "source": {
                "path": args.source.as_posix(),
                "architecture": args.source_arch,
            },
            "s3": {
                "bucket_name": args.bucket,
            },
            "images": {
                source_sha256: {
                    "boot_mode": args.boot_mode,
                    "regions": [
                        args.region,
                    ],
                    "public": False,
                    "description": "a test image",
                    "share": args.share,
                    "root_device_volume_size": args.root_device_volume_size,
                }
            },
        }
    }
    return config


def _parser():
    parser = argparse.ArgumentParser(description="Publish a file on AWS EC2")
    parser.add_argument("--log-level", choices=["info", "debug"], default="info")
    parser.add_argument(
        "--source-arch",
        type=str,
        choices=["x86_64", "arm64"],
        default="x86_64",
        help="architecture of the source file",
    )
    parser.add_argument(
        "--boot-mode",
        type=str,
        choices=["legacy-bios", "uefi-preferred", "uefi"],
        default="uefi",
        help="the boot mode to use. default: %(default)s",
    )
    parser.add_argument(
        "--region",
        type=str,
        default="eu-central-1",
        help="the region to register the image in. default: %(default)s",
    )
    parser.add_argument(
        "--root-device-volume-size",
        type=int,
        default=20,
        help="root device volume size. default: %(default)s G",
    )
    parser.add_argument(
        "--share", nargs="+", default=[], help="list of accounts to share with"
    )
    parser.add_argument(
        "source", type=pathlib.Path, help="file to publish (a .vmdk file)"
    )
    parser.add_argument("bucket", type=str, help="a S3 bucket name to use")
    return parser


def main():
    parser = _parser()
    args = parser.parse_args()

    log_formatter = logging.Formatter("%(asctime)s:%(name)s:%(levelname)s:%(message)s")
    # log level
    loglevel = logging.INFO
    if args.log_level == "debug":
        loglevel = logging.DEBUG
    root_logger = logging.getLogger()
    root_logger.setLevel(loglevel)
    # log console
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(log_formatter)
    root_logger.addHandler(console_handler)

    with tempfile.TemporaryDirectory() as tmpdir:
        config_path = pathlib.Path(tmpdir) / "config.yaml"
        with open(config_path, "w") as f:
            yaml.dump(_config_create(args), f)
        images_by_name, images_by_group = awspub.create(config_path, None, None)
        print(
            "----- YOU CAN RUN THE FOLLOWING COMMAND(S) TO CREATE THE INSTANCES -----"
        )
        if args.source_arch == "x86_64":
            instance_type = "m6a.large"
        else:
            instance_type = "m6g.large"
        for image_name, values in images_by_name.items():
            for region, ami_id in values.items():
                print(
                    f"aws ec2 run-instances --image-id {ami_id} --instance-type {instance_type} --region {region} --key-name toabctl --tag-specifications 'ResourceType=instance,Tags=[{{Key=Name,Value=toabctl-awspub-register-{image_name}}}]'"
                )


if __name__ == "__main__":
    main()