aboutsummaryrefslogtreecommitdiff
path: root/contrib/cloud/aws-import
blob: a93a00843bfd6537ac7f2d6736e070a177cc174e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
#!/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


def launch_link(region, image_id):
    """Construct a web console launch link"""
    return ("https://console.aws.amazon.com/ec2/v2/home?"
            "region=%s#LaunchInstanceWizard:ami=%s" % (region, 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('--wiki', '-w', metavar='FILE',
                    help="Generate Dokuwiki table")
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)}

# Construct Dokuwiki table
wikitab = ["^ AWS region  ^ CPU architecture  ^ AMI ID  ^\n"] + list(
    "| ''%s''  | ''%s''  | ''[[%s|%s]]''  |\n" % (
        region,
        args.architecture,
        launch_link(region, results[region]),
        results[region],
    ) for region in args.region)
if args.wiki:
    with open(args.wiki, 'wt') as fh:
        fh.writelines(wikitab)

# Show created images
for region in args.region:
    print("%s: %s" % (region, results[region]))