Automate customization of an existing AMI

modify-ami.py

#!/bin/env python
#
# Automate customization of an existing AMI. Intended for producing jenkins
# slaves for multiple platforms in easily reproducable / customizable fashion.
#
# Runs following steps unattended:
# 1) create a new instance (m1.small) from given AMI
# 2) wait for it to get started and SSH accessible
# 3) run a set of commands on the instance
# 4) stop the instance (and wait for it to stop)
# 5) make a snapshot of the root filesystem
# 6) terminate the instance
#
# Some day maybe:
# 7) wait for the snapshot to complete
# 8) create a new image from the snapshot (making sure to set right kernel)
#
# Notes:
# - ssh-add your ssh private key before you begin
# - you can probably get an interactive shell instead/in addition to fixed
#   script
# - AMIs have constraints on which instance types they can run, but there
#   doesn't seem to be any way to search for compatible ones
#
import subprocess
import time

import boto


API_KEY = SETME
API_SECRET = SETME
KEY_NAME = SETME


def start_instances(ami_id):

    conn = boto.connect_ec2(aws_access_key_id=API_KEY,
                            aws_secret_access_key=API_SECRET,
                            )
    ami = conn.get_image(ami_id)
    print(ami)
    # print(ami.__dict__)
    assert ami.root_device_type == 'ebs', ami.root_device_name

    reservation = ami.run(instance_type='m1.small',
                          key_name=KEY_NAME,
                          )
    print('{0.id}: {0.instances}'.format(reservation))

    tags = {'Name': 'modify-ami instance of %s' % ami.id,
            'Role': 'throwaway',
            }
    for instance in reservation.instances:
        instance.update()
    conn.create_tags([instance.id for instance in reservation.instances],
                     tags)

    wait_for_state('running', reservation.instances)

    for instance in reservation.instances:
        print('{0.id}: ssh ubuntu@{0.public_dns_name}'.format(instance))

    return reservation.instances


def wait_for_state(states, instances):
    if isinstance(states, basestring):
        states = [states]

    pending_instances = instances[:]
    while pending_instances:
        for instance in pending_instances:
            instance.update()  # refresh local status
            print '{0.id}: {0.state}'.format(instance)
            if instance.state in states:
                pending_instances.remove(instance)
        time.sleep(2)


def install_java(
        instance=None,
        hostname=None,
        target='/opt/java-for-jenkins',
        source='http://javadl.sun.com/webapps/download/AutoDL?BundleId=106240',
        filename='jre1.8.0_45.tgz',
        ):
    if hostname is None:
        hostname = instance.public_dns_name

    install_script = '''\
curl -L "{source}" -o /tmp/{filename}
sudo mkdir -p {target}
sudo tar xzf /tmp/{filename} -C {target} --strip-components=1
rm -f /tmp/{filename}'''.format(**locals())

    s = ssh(hostname, 'bash', '-xe', stdin=subprocess.PIPE)
    s.stdin.write(install_script)
    s.stdin.close()
    s.communicate()


def stop_instances(instances):
    for instance in instances:
        instance.stop()

    wait_for_state(['stopped', 'terminated'], instances)


def ssh(hostname, *args, **kwargs):
    username = kwargs.pop('username', 'ubuntu')
    cmd = ('ssh',
           '-o', 'UserKnownHostsFile=/dev/null',
           '-o', 'StrictHostKeyChecking=no',
           '{username}@{hostname}'.format(**locals()),
           ) + args
    return subprocess.Popen(cmd, **kwargs)


def wait_for_ssh(instance):
    while ssh(instance.public_dns_name, 'true').wait():
        time.sleep(1)
    print '{0.id}: ssh up'.format(instance)


if __name__ == '__main__':
    # ami_id = 'ami-59a4a230'

    instances = start_instances(ami_id)

    for instance in instances:
        wait_for_ssh(instance)
        install_java(instance=instance)

    stop_instances(instances)

    for instance in instances:
        mapping = instance.block_device_mapping[instance.root_device_name]
        print('{instance.id}: volume_id={mapping.volume_id}'.format(**locals()))

        volume = instance.connection.get_all_volumes([mapping.volume_id])[0]

        snapshot = volume.create_snapshot(
            description='{ami_id} with java-for-jenkins'.format(ami_id=ami_id),
        )
        print('{0.id}: {0.status}'.format(snapshot.id))

        instance.terminate()

    wait_for_state('terminated', instances)

    # wait_for_state('done?', snapshots)
    # conn.create_image(snapshot, ...)  # need to fix kernel?!

Could probably just pass the customization script in instance.run(user_data='') and save steps.