upload-ami

#!/usr/bin/python3 -u

# THIS SCRIPT SHOULD BE CONSIDERED DEPRECATED; please avoid extending it!
# Instead, add new functionality into coreos-assembler.
#
# This wraps coreos-assembler buildextend-aws which itself wraps `ore aws upload`,
# but also supports making the images public because ore doesn't do that right
# now.  For FCOS, that's handled via plume, but RHCOS as of right now doesn't
# use that.  But, what would make sense here is to move some of this functionality
# into cosa buildextend-aws, such as the marketplace support.

import argparse
import json
import os
import subprocess
import sys

OPENSHIFT_CI_ACCOUNT = "460538899914"
OPENSHIFT_DEV_ACCOUNT = "269733383066"
OPENSHIFT_MARKETPLACE_TEST_ACCOUNT = "125666959065"

# Since we run all tests on QEMU, and as of today don't have really
# any tests that are specific to a cloud provider, we just run
# basic sanity checks.
SANITY_CHECK = "rhcos.basic"
# Don't need anything big
INSTANCE_TYPE = 't2.small'

AWS = 'aws'
ALIYUN = 'aliyun'

cloud = None


# Upload the given build to a bucket in a region
def upload_to_cloud(buildid, region, bucket, suffix=None, public=False):
    print("Uploading {} to {} in {}...".format(buildid, bucket, region))
    ca_args = ['coreos-assembler',
               f'buildextend-{cloud}',
               '--upload',
               f'--build={buildid}',
               f'--region={region}',
               f'--bucket={bucket}']

    if suffix is not None:
        ca_args.append(f'--name-suffix={suffix}')

    # If the AMI is public, grant image launch and ec2 snapshot permissions to Marketplace account.
    if public and cloud == AWS:
        ca_args.extend(["--grant-user", OPENSHIFT_MARKETPLACE_TEST_ACCOUNT])
        ca_args.extend(["--grant-user-snapshot", OPENSHIFT_MARKETPLACE_TEST_ACCOUNT])
    try:
        print(ca_args)
        subprocess.check_call(ca_args)
    except subprocess.CalledProcessError as upload_err:
        print("Failed to build/upload image! Error: {}".format(upload_err))
        raise

    # If the AMI isn't public, then just grant access to the dev and CI accounts.
    if not public and cloud == AWS:
        ca_args.extend(["--grant-user", OPENSHIFT_CI_ACCOUNT, OPENSHIFT_DEV_ACCOUNT])
    try:
        print(ca_args)
        subprocess.check_call(ca_args)
    except subprocess.CalledProcessError as upload_err:
        print("Failed to build/upload image! Error: {}".format(upload_err))
        raise

# Run the basic kola test on an AMI
def run_kola_basic(ami_id, region):
    kola_args = ['kola', '-b', 'rhcos',
                 '-p', 'aws',
                 '--ignition-version', 'v2',
                 '--aws-type', INSTANCE_TYPE,
                 '--tapfile', 'rhcos-aws.tap',
                 '--aws-ami', ami_id,
                 '--aws-region', region,
                 'run', SANITY_CHECK]

    try:
        subprocess.check_call(kola_args)
    except subprocess.CalledProcessError as kola_err:
        print("Running kola tests on AWS failed!  Error: {}".format(kola_err))
        raise

# Set all the permissions to make the AMI/snapshot public
def make_ami_public(region, ami_id):
    print(f"Making {ami_id} public")
    # allow the AMI to started by everyone
    try:
        subprocess.check_call(['aws', 'ec2', "--region", region,
                               'modify-image-attribute',
                               '--image-id', ami_id,
                               '--launch-permission',
                               '{"Add": [{"Group":"all"}]}'])
    except subprocess.CalledProcessError as pub_err:
        print("Failed to make {} public! Error: {}".format(ami_id, pub_err))
        raise

    # check that the AMI is actually public
    try:
        output = subprocess.check_output(['aws', 'ec2',
                                          '--region', region,
                                          'describe-images',
                                          '--image-id', ami_id])
    except subprocess.CalledProcessError as pub_err:
        print("Failed to describe AMI {}! Error: {}".format(ami_id, pub_err))
        raise

    image_description = json.loads(output)
    snapshot = None
    for image in image_description['Images']:
        for block_device_mapping in image['BlockDeviceMappings']:
            ebs = block_device_mapping.get('Ebs')
            if ebs is not None:
                snapshot = ebs['SnapshotId']
                break

    assert snapshot is not None, \
        "Failed to find SnapshotId associated with AMI {}".format(ami_id)

    # allow everyone to create volumes with the snapshot
    print("Adding createVolumePermission for snapshot {}".format(snapshot) +
          "in region {}".format(region))
    try:
        subprocess.check_call(['aws', 'ec2', "--region", region,
                               'modify-snapshot-attribute',
                               '--snapshot-id', snapshot,
                               '--attribute', 'createVolumePermission',
                               '--operation-type', 'add',
                               '--group-names', 'all'])
    except subprocess.CalledProcessError as pub_err:
        print("Failed to add createVolumePermission on " +
              "snapshot {}!  Error: {}".format(snapshot, pub_err))
        raise


def main():
    # the aliyun API is so close to the AWS one, we're just using the same
    # script for both of them
    entrypoint = os.path.basename(sys.argv[0])
    global cloud
    if entrypoint == 'upload-ami':
        cloud = AWS
    elif entrypoint == 'upload-aliyun-image':
        cloud = ALIYUN
    else:
        assert False, f"Bad entrypoint {entrypoint}"

    # parse args and dispatch
    parser = argparse.ArgumentParser()
    parser.add_argument("--build", action='store', required=True)
    parser.add_argument("--region", action='store', required=True)
    parser.add_argument("--bucket", action='store', required=True)
    parser.add_argument("--name-suffix", action='store')
    parser.add_argument("--public", action='store_true')
    parser.add_argument("--skip-kola", action='store_true')
    args = parser.parse_args()

    # we don't support this flag on aliyun yet, but we still want it in the
    # args namespace
    if cloud == ALIYUN:
        assert not args.public

    try:
        upload_to_cloud(args.build, args.region,
                        args.bucket, args.name_suffix,
                        args.public)
    except subprocess.CalledProcessError:
        sys.exit(1)

    arch = subprocess.check_output(["cosa", "basearch"], text='UTF-8').strip()

    if args.public:
        with open(f"builds/{args.build}/{arch}/meta.json") as meta_f:
            meta = json.load(meta_f)
        ami_id = None
        for region in meta['amis']:
            if region['name'] == args.region:
                ami_id = region['hvm']
        if ami_id is None:
            print("Failed to find AMI {} in meta.json!".format(ami_id))
            sys.exit(1)

        if not args.skip_kola:
            try:
                run_kola_basic(ami_id, args.region)
            except subprocess.CalledProcessError:
                sys.exit(1)

        # Actually make the AMI/Snapshot public now
        try:
            make_ami_public(args.region, ami_id)
        except (AssertionError, subprocess.CalledProcessError):
            sys.exit(1)


if __name__ == '__main__':
    main()