Merge branch 'master' into openslxopenslx

1 files changed, 53 insertions, 14 deletions

diff --git a/contrib/cloud/aws-import b/contrib/cloud/aws-import
index ace005870..77c0fd0f7 100755
--- a/contrib/cloud/aws-import
+++ b/contrib/cloud/aws-import
@@ -22,11 +22,12 @@ def detect_architecture(image):
     return 'x86_64'
 
 
-def create_snapshot(region, description, image):
+def create_snapshot(region, description, image, tags):
     """Create an EBS snapshot"""
     client = boto3.client('ebs', region_name=region)
     snapshot = client.start_snapshot(VolumeSize=1,
-                                     Description=description)
+                                     Description=description,
+                                     Tags=tags)
     snapshot_id = snapshot['SnapshotId']
     with open(image, 'rb') as fh:
         for block in count():
@@ -46,21 +47,42 @@ def create_snapshot(region, description, image):
     return snapshot_id
 
 
-def import_image(region, name, architecture, image, public, overwrite):
-    """Import an AMI image"""
+def delete_images(region, filters, retain):
     client = boto3.client('ec2', region_name=region)
     resource = boto3.resource('ec2', region_name=region)
-    description = '%s (%s)' % (name, architecture)
-    images = client.describe_images(Filters=[{'Name': 'name',
-                                              'Values': [description]}])
-    if overwrite and images['Images']:
-        images = images['Images'][0]
-        image_id = images['ImageId']
-        snapshot_id = images['BlockDeviceMappings'][0]['Ebs']['SnapshotId']
+    images = client.describe_images(Owners=['self'], Filters=filters)
+    old_images = sorted(images['Images'], key=lambda x: x['CreationDate'])
+    if retain > 0:
+        old_images = old_images[:-retain]
+    for image in old_images:
+        image_id = image['ImageId']
+        snapshot_id = image['BlockDeviceMappings'][0]['Ebs']['SnapshotId']
         resource.Image(image_id).deregister()
         resource.Snapshot(snapshot_id).delete()
+
+
+def import_image(region, name, family, architecture, image, public, overwrite,
+                 retain):
+    """Import an AMI image"""
+    client = boto3.client('ec2', region_name=region)
+    resource = boto3.resource('ec2', region_name=region)
+    description = '%s (%s)' % (name, architecture)
+    tags = [
+        {'Key': 'family', 'Value': family},
+        {'Key': 'architecture', 'Value': architecture},
+    ]
+    if overwrite:
+        filters = [{'Name': 'name', 'Values': [description]}]
+        delete_images(region=region, filters=filters, retain=0)
+    if retain is not None:
+        filters = [
+            {'Name': 'tag:family', 'Values': [family]},
+            {'Name': 'tag:architecture', 'Values': [architecture]},
+            {'Name': 'is-public', 'Values': [str(public).lower()]},
+        ]
+        delete_images(region=region, filters=filters, retain=retain)
     snapshot_id = create_snapshot(region=region, description=description,
-                                  image=image)
+                                  image=image, tags=tags)
     client.get_waiter('snapshot_completed').wait(SnapshotIds=[snapshot_id])
     image = client.register_image(Architecture=architecture,
                                   BlockDeviceMappings=[{
@@ -72,12 +94,19 @@ def import_image(region, name, architecture, image, public, overwrite):
                                   }],
                                   EnaSupport=True,
                                   Name=description,
+                                  TagSpecifications=[{
+                                      'ResourceType': 'image',
+                                      'Tags': tags,
+                                  }],
                                   RootDeviceName='/dev/sda1',
                                   SriovNetSupport='simple',
                                   VirtualizationType='hvm')
     image_id = image['ImageId']
     client.get_waiter('image_available').wait(ImageIds=[image_id])
     if public:
+        image_block = client.get_image_block_public_access_state()
+        if image_block['ImageBlockPublicAccessState'] != 'unblocked':
+            client.disable_image_block_public_access()
         resource.Image(image_id).modify_attribute(Attribute='launchPermission',
                                                   OperationType='add',
                                                   UserGroups=['all'])
@@ -94,10 +123,14 @@ def launch_link(region, image_id):
 parser = argparse.ArgumentParser(description="Import AWS EC2 image (AMI)")
 parser.add_argument('--name', '-n',
                     help="Image name")
+parser.add_argument('--family', '-f',
+                    help="Image family name")
 parser.add_argument('--public', '-p', action='store_true',
                     help="Make image public")
 parser.add_argument('--overwrite', action='store_true',
                     help="Overwrite any existing image with same name")
+parser.add_argument('--retain', type=int, metavar='NUM',
+                    help="Retain at most <NUM> old images")
 parser.add_argument('--region', '-r', action='append',
                     help="AWS region(s)")
 parser.add_argument('--wiki', '-w', metavar='FILE',
@@ -108,9 +141,13 @@ args = parser.parse_args()
 # Detect CPU architectures
 architectures = {image: detect_architecture(image) for image in args.image}
 
+# Use default family name if none specified
+if not args.family:
+    args.family = 'iPXE'
+
 # Use default name if none specified
 if not args.name:
-    args.name = 'iPXE (%s)' % date.today().strftime('%Y-%m-%d')
+    args.name = '%s (%s)' % (args.family, date.today().strftime('%Y-%m-%d'))
 
 # Use all regions if none specified
 if not args.region:
@@ -123,10 +160,12 @@ with ThreadPoolExecutor(max_workers=len(imports)) as executor:
     futures = {executor.submit(import_image,
                                region=region,
                                name=args.name,
+                               family=args.family,
                                architecture=architectures[image],
                                image=image,
                                public=args.public,
-                               overwrite=args.overwrite): (region, image)
+                               overwrite=args.overwrite,
+                               retain=args.retain): (region, image)
                for region, image in imports}
     results = {futures[future]: future.result()
                for future in as_completed(futures)}