blob: ace005870398d5df8f5e51a2d30e71d53ab304a2 [file] [log] [blame]
Michael Brownd16535a2021-02-16 00:27:40 +00001#!/usr/bin/env python3
2
3import argparse
4from base64 import b64encode
5from concurrent.futures import ThreadPoolExecutor, as_completed
Michael Brown6dad3162021-05-01 22:08:17 +01006from datetime import date
Michael Brownd16535a2021-02-16 00:27:40 +00007from hashlib import sha256
8from itertools import count
Michael Brown438513f2021-05-02 09:39:10 +01009import subprocess
Michael Brownd16535a2021-02-16 00:27:40 +000010
11import boto3
12
13BLOCKSIZE = 512 * 1024
14
15
Michael Brown438513f2021-05-02 09:39:10 +010016def detect_architecture(image):
17 """Detect CPU architecture"""
18 mdir = subprocess.run(['mdir', '-b', '-i', image, '::/EFI/BOOT'],
Michael Brown70995392022-04-06 14:36:07 +010019 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
Michael Brown438513f2021-05-02 09:39:10 +010020 if any(b'BOOTAA64.EFI' in x for x in mdir.stdout.splitlines()):
21 return 'arm64'
22 return 'x86_64'
23
24
Michael Brownd16535a2021-02-16 00:27:40 +000025def create_snapshot(region, description, image):
26 """Create an EBS snapshot"""
27 client = boto3.client('ebs', region_name=region)
28 snapshot = client.start_snapshot(VolumeSize=1,
29 Description=description)
30 snapshot_id = snapshot['SnapshotId']
31 with open(image, 'rb') as fh:
32 for block in count():
33 data = fh.read(BLOCKSIZE)
34 if not data:
35 break
36 data = data.ljust(BLOCKSIZE, b'\0')
37 checksum = b64encode(sha256(data).digest()).decode()
38 client.put_snapshot_block(SnapshotId=snapshot_id,
39 BlockIndex=block,
40 BlockData=data,
41 DataLength=BLOCKSIZE,
42 Checksum=checksum,
43 ChecksumAlgorithm='SHA256')
44 client.complete_snapshot(SnapshotId=snapshot_id,
45 ChangedBlocksCount=block)
46 return snapshot_id
47
48
Michael Brownd8f9c222023-11-07 15:54:59 +000049def import_image(region, name, architecture, image, public, overwrite):
Michael Brownd16535a2021-02-16 00:27:40 +000050 """Import an AMI image"""
51 client = boto3.client('ec2', region_name=region)
52 resource = boto3.resource('ec2', region_name=region)
53 description = '%s (%s)' % (name, architecture)
Michael Brownd8f9c222023-11-07 15:54:59 +000054 images = client.describe_images(Filters=[{'Name': 'name',
55 'Values': [description]}])
56 if overwrite and images['Images']:
57 images = images['Images'][0]
58 image_id = images['ImageId']
59 snapshot_id = images['BlockDeviceMappings'][0]['Ebs']['SnapshotId']
60 resource.Image(image_id).deregister()
61 resource.Snapshot(snapshot_id).delete()
Michael Brownd16535a2021-02-16 00:27:40 +000062 snapshot_id = create_snapshot(region=region, description=description,
63 image=image)
64 client.get_waiter('snapshot_completed').wait(SnapshotIds=[snapshot_id])
65 image = client.register_image(Architecture=architecture,
66 BlockDeviceMappings=[{
67 'DeviceName': '/dev/sda1',
68 'Ebs': {
69 'SnapshotId': snapshot_id,
70 'VolumeType': 'standard',
71 },
72 }],
73 EnaSupport=True,
74 Name=description,
75 RootDeviceName='/dev/sda1',
76 SriovNetSupport='simple',
77 VirtualizationType='hvm')
78 image_id = image['ImageId']
79 client.get_waiter('image_available').wait(ImageIds=[image_id])
80 if public:
81 resource.Image(image_id).modify_attribute(Attribute='launchPermission',
82 OperationType='add',
83 UserGroups=['all'])
84 return image_id
85
86
Michael Browne9942372021-05-01 21:33:38 +010087def launch_link(region, image_id):
88 """Construct a web console launch link"""
89 return ("https://console.aws.amazon.com/ec2/v2/home?"
90 "region=%s#LaunchInstanceWizard:ami=%s" % (region, image_id))
91
92
Michael Brownd16535a2021-02-16 00:27:40 +000093# Parse command-line arguments
94parser = argparse.ArgumentParser(description="Import AWS EC2 image (AMI)")
Michael Brown6dad3162021-05-01 22:08:17 +010095parser.add_argument('--name', '-n',
Michael Brownd16535a2021-02-16 00:27:40 +000096 help="Image name")
97parser.add_argument('--public', '-p', action='store_true',
98 help="Make image public")
Michael Brownd8f9c222023-11-07 15:54:59 +000099parser.add_argument('--overwrite', action='store_true',
100 help="Overwrite any existing image with same name")
Michael Brownd16535a2021-02-16 00:27:40 +0000101parser.add_argument('--region', '-r', action='append',
102 help="AWS region(s)")
Michael Browne9942372021-05-01 21:33:38 +0100103parser.add_argument('--wiki', '-w', metavar='FILE',
104 help="Generate Dokuwiki table")
Michael Brown106f4c52021-05-02 12:23:00 +0100105parser.add_argument('image', nargs='+', help="iPXE disk image")
Michael Brownd16535a2021-02-16 00:27:40 +0000106args = parser.parse_args()
107
Michael Brown106f4c52021-05-02 12:23:00 +0100108# Detect CPU architectures
109architectures = {image: detect_architecture(image) for image in args.image}
Michael Brown438513f2021-05-02 09:39:10 +0100110
Michael Brown6dad3162021-05-01 22:08:17 +0100111# Use default name if none specified
112if not args.name:
Michael Brown106f4c52021-05-02 12:23:00 +0100113 args.name = 'iPXE (%s)' % date.today().strftime('%Y-%m-%d')
Michael Brown6dad3162021-05-01 22:08:17 +0100114
Michael Brownd16535a2021-02-16 00:27:40 +0000115# Use all regions if none specified
116if not args.region:
117 args.region = sorted(x['RegionName'] for x in
118 boto3.client('ec2').describe_regions()['Regions'])
119
Michael Brown106f4c52021-05-02 12:23:00 +0100120# Use one thread per import to maximise parallelism
121imports = [(region, image) for region in args.region for image in args.image]
122with ThreadPoolExecutor(max_workers=len(imports)) as executor:
Michael Brownd16535a2021-02-16 00:27:40 +0000123 futures = {executor.submit(import_image,
124 region=region,
125 name=args.name,
Michael Brown106f4c52021-05-02 12:23:00 +0100126 architecture=architectures[image],
127 image=image,
Michael Brownd8f9c222023-11-07 15:54:59 +0000128 public=args.public,
129 overwrite=args.overwrite): (region, image)
Michael Brown106f4c52021-05-02 12:23:00 +0100130 for region, image in imports}
Michael Brownd16535a2021-02-16 00:27:40 +0000131 results = {futures[future]: future.result()
132 for future in as_completed(futures)}
133
Michael Browne9942372021-05-01 21:33:38 +0100134# Construct Dokuwiki table
135wikitab = ["^ AWS region ^ CPU architecture ^ AMI ID ^\n"] + list(
136 "| ''%s'' | ''%s'' | ''[[%s|%s]]'' |\n" % (
137 region,
Michael Brown106f4c52021-05-02 12:23:00 +0100138 architectures[image],
139 launch_link(region, results[(region, image)]),
140 results[(region, image)],
141 ) for region, image in imports)
Michael Browne9942372021-05-01 21:33:38 +0100142if args.wiki:
143 with open(args.wiki, 'wt') as fh:
144 fh.writelines(wikitab)
145
Michael Brownd16535a2021-02-16 00:27:40 +0000146# Show created images
Michael Brown106f4c52021-05-02 12:23:00 +0100147for region, image in imports:
148 print("%s %s %s %s" % (
149 region, image, architectures[image], results[(region, image)]
150 ))