[cloud] Add utility for importing images to AWS EC2

Add a utility that can be used to upload an iPXE disk image to AWS EC2
as an Amazon Machine Image (AMI).  For example:

  make CONFIG=cloud EMBED=config/cloud/aws.ipxe bin/ipxe.usb

  ../contrib/cloud/aws-import -p -n "iPXE 1.21.1" bin/ipxe.usb

Uploads are performed in parallel across all regions, and use the EBS
direct APIs to avoid the need to store temporary files in S3 or to run
VM import tasks.

Signed-off-by: Michael Brown <mcb30@ipxe.org>

1 files changed, 100 insertions, 0 deletions

diff --git a/contrib/cloud/aws-import b/contrib/cloud/aws-import
new file mode 100755
index 00000000..9ee53e70
--- /dev/null
+++ b/contrib/cloud/aws-import
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+
+import argparse
+from base64 import b64encode
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from hashlib import sha256
+from itertools import count
+
+import boto3
+
+BLOCKSIZE = 512 * 1024
+
+
+def create_snapshot(region, description, image):
+    """Create an EBS snapshot"""
+    client = boto3.client('ebs', region_name=region)
+    snapshot = client.start_snapshot(VolumeSize=1,
+                                     Description=description)
+    snapshot_id = snapshot['SnapshotId']
+    with open(image, 'rb') as fh:
+        for block in count():
+            data = fh.read(BLOCKSIZE)
+            if not data:
+                break
+            data = data.ljust(BLOCKSIZE, b'\0')
+            checksum = b64encode(sha256(data).digest()).decode()
+            client.put_snapshot_block(SnapshotId=snapshot_id,
+                                      BlockIndex=block,
+                                      BlockData=data,
+                                      DataLength=BLOCKSIZE,
+                                      Checksum=checksum,
+                                      ChecksumAlgorithm='SHA256')
+    client.complete_snapshot(SnapshotId=snapshot_id,
+                             ChangedBlocksCount=block)
+    return snapshot_id
+
+
+def import_image(region, name, architecture, image, public):
+    """Import an AMI image"""
+    client = boto3.client('ec2', region_name=region)
+    resource = boto3.resource('ec2', region_name=region)
+    description = '%s (%s)' % (name, architecture)
+    snapshot_id = create_snapshot(region=region, description=description,
+                                  image=image)
+    client.get_waiter('snapshot_completed').wait(SnapshotIds=[snapshot_id])
+    image = client.register_image(Architecture=architecture,
+                                  BlockDeviceMappings=[{
+                                      'DeviceName': '/dev/sda1',
+                                      'Ebs': {
+                                          'SnapshotId': snapshot_id,
+                                          'VolumeType': 'standard',
+                                      },
+                                  }],
+                                  EnaSupport=True,
+                                  Name=description,
+                                  RootDeviceName='/dev/sda1',
+                                  SriovNetSupport='simple',
+                                  VirtualizationType='hvm')
+    image_id = image['ImageId']
+    client.get_waiter('image_available').wait(ImageIds=[image_id])
+    if public:
+        resource.Image(image_id).modify_attribute(Attribute='launchPermission',
+                                                  OperationType='add',
+                                                  UserGroups=['all'])
+    return image_id
+
+
+# Parse command-line arguments
+parser = argparse.ArgumentParser(description="Import AWS EC2 image (AMI)")
+parser.add_argument('--architecture', '-a', default='x86_64',
+                    help="CPU architecture")
+parser.add_argument('--name', '-n', required=True,
+                    help="Image name")
+parser.add_argument('--public', '-p', action='store_true',
+                    help="Make image public")
+parser.add_argument('--region', '-r', action='append',
+                    help="AWS region(s)")
+parser.add_argument('image', help="iPXE disk image")
+args = parser.parse_args()
+
+# Use all regions if none specified
+if not args.region:
+    args.region = sorted(x['RegionName'] for x in
+                         boto3.client('ec2').describe_regions()['Regions'])
+
+# Use one thread per region to maximise parallelism
+with ThreadPoolExecutor(max_workers=len(args.region)) as executor:
+    futures = {executor.submit(import_image,
+                               region=region,
+                               name=args.name,
+                               architecture=args.architecture,
+                               image=args.image,
+                               public=args.public): region
+               for region in args.region}
+    results = {futures[future]: future.result()
+               for future in as_completed(futures)}
+
+# Show created images
+for region in args.region:
+    print("%s: %s" % (region, results[region]))