Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
| 2 | |
| 3 | import argparse |
| 4 | from base64 import b64encode |
| 5 | from concurrent.futures import ThreadPoolExecutor, as_completed |
Michael Brown | 6dad316 | 2021-05-01 22:08:17 +0100 | [diff] [blame] | 6 | from datetime import date |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 7 | from hashlib import sha256 |
| 8 | from itertools import count |
Michael Brown | 438513f | 2021-05-02 09:39:10 +0100 | [diff] [blame] | 9 | import subprocess |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 10 | |
| 11 | import boto3 |
| 12 | |
| 13 | BLOCKSIZE = 512 * 1024 |
| 14 | |
| 15 | |
Michael Brown | 438513f | 2021-05-02 09:39:10 +0100 | [diff] [blame] | 16 | def detect_architecture(image): |
| 17 | """Detect CPU architecture""" |
| 18 | mdir = subprocess.run(['mdir', '-b', '-i', image, '::/EFI/BOOT'], |
Michael Brown | 7099539 | 2022-04-06 14:36:07 +0100 | [diff] [blame] | 19 | stdout=subprocess.PIPE, stderr=subprocess.PIPE) |
Michael Brown | 438513f | 2021-05-02 09:39:10 +0100 | [diff] [blame] | 20 | if any(b'BOOTAA64.EFI' in x for x in mdir.stdout.splitlines()): |
| 21 | return 'arm64' |
| 22 | return 'x86_64' |
| 23 | |
| 24 | |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 25 | def 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 Brown | d8f9c22 | 2023-11-07 15:54:59 +0000 | [diff] [blame] | 49 | def import_image(region, name, architecture, image, public, overwrite): |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 50 | """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 Brown | d8f9c22 | 2023-11-07 15:54:59 +0000 | [diff] [blame] | 54 | 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 Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 62 | 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 Brown | e994237 | 2021-05-01 21:33:38 +0100 | [diff] [blame] | 87 | def 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 Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 93 | # Parse command-line arguments |
| 94 | parser = argparse.ArgumentParser(description="Import AWS EC2 image (AMI)") |
Michael Brown | 6dad316 | 2021-05-01 22:08:17 +0100 | [diff] [blame] | 95 | parser.add_argument('--name', '-n', |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 96 | help="Image name") |
| 97 | parser.add_argument('--public', '-p', action='store_true', |
| 98 | help="Make image public") |
Michael Brown | d8f9c22 | 2023-11-07 15:54:59 +0000 | [diff] [blame] | 99 | parser.add_argument('--overwrite', action='store_true', |
| 100 | help="Overwrite any existing image with same name") |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 101 | parser.add_argument('--region', '-r', action='append', |
| 102 | help="AWS region(s)") |
Michael Brown | e994237 | 2021-05-01 21:33:38 +0100 | [diff] [blame] | 103 | parser.add_argument('--wiki', '-w', metavar='FILE', |
| 104 | help="Generate Dokuwiki table") |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 105 | parser.add_argument('image', nargs='+', help="iPXE disk image") |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 106 | args = parser.parse_args() |
| 107 | |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 108 | # Detect CPU architectures |
| 109 | architectures = {image: detect_architecture(image) for image in args.image} |
Michael Brown | 438513f | 2021-05-02 09:39:10 +0100 | [diff] [blame] | 110 | |
Michael Brown | 6dad316 | 2021-05-01 22:08:17 +0100 | [diff] [blame] | 111 | # Use default name if none specified |
| 112 | if not args.name: |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 113 | args.name = 'iPXE (%s)' % date.today().strftime('%Y-%m-%d') |
Michael Brown | 6dad316 | 2021-05-01 22:08:17 +0100 | [diff] [blame] | 114 | |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 115 | # Use all regions if none specified |
| 116 | if not args.region: |
| 117 | args.region = sorted(x['RegionName'] for x in |
| 118 | boto3.client('ec2').describe_regions()['Regions']) |
| 119 | |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 120 | # Use one thread per import to maximise parallelism |
| 121 | imports = [(region, image) for region in args.region for image in args.image] |
| 122 | with ThreadPoolExecutor(max_workers=len(imports)) as executor: |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 123 | futures = {executor.submit(import_image, |
| 124 | region=region, |
| 125 | name=args.name, |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 126 | architecture=architectures[image], |
| 127 | image=image, |
Michael Brown | d8f9c22 | 2023-11-07 15:54:59 +0000 | [diff] [blame] | 128 | public=args.public, |
| 129 | overwrite=args.overwrite): (region, image) |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 130 | for region, image in imports} |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 131 | results = {futures[future]: future.result() |
| 132 | for future in as_completed(futures)} |
| 133 | |
Michael Brown | e994237 | 2021-05-01 21:33:38 +0100 | [diff] [blame] | 134 | # Construct Dokuwiki table |
| 135 | wikitab = ["^ AWS region ^ CPU architecture ^ AMI ID ^\n"] + list( |
| 136 | "| ''%s'' | ''%s'' | ''[[%s|%s]]'' |\n" % ( |
| 137 | region, |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 138 | architectures[image], |
| 139 | launch_link(region, results[(region, image)]), |
| 140 | results[(region, image)], |
| 141 | ) for region, image in imports) |
Michael Brown | e994237 | 2021-05-01 21:33:38 +0100 | [diff] [blame] | 142 | if args.wiki: |
| 143 | with open(args.wiki, 'wt') as fh: |
| 144 | fh.writelines(wikitab) |
| 145 | |
Michael Brown | d16535a | 2021-02-16 00:27:40 +0000 | [diff] [blame] | 146 | # Show created images |
Michael Brown | 106f4c5 | 2021-05-02 12:23:00 +0100 | [diff] [blame] | 147 | for region, image in imports: |
| 148 | print("%s %s %s %s" % ( |
| 149 | region, image, architectures[image], results[(region, image)] |
| 150 | )) |