From d16535aa4fcb30c78559103e1bc994fa99bf5748 Mon Sep 17 00:00:00 2001
From: Michael Brown <mcb30@ipxe.org>
Date: Tue, 16 Feb 2021 00:27:40 +0000
Subject: [PATCH] [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>
---
 contrib/cloud/aws-import | 100 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 100 insertions(+)
 create mode 100755 contrib/cloud/aws-import

diff --git a/contrib/cloud/aws-import b/contrib/cloud/aws-import
new file mode 100755
index 000000000..9ee53e704
--- /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]))