Flask: Mock EC2 Endpoints

ec2_endpoints.py

from flask import Flask
import re

app = Flask(__name__)
script = '''#!/bin/bash

cat <<EOF
============================
My name is ${0}
I was input via user data
============================
EOF

touch /var/log/done
'''

meta = {
    'user-data': script,
    'meta-data': {
        'ami-id': '00000000',
        'ami-launch-index': '0',
        'ami-manifest-path': '',
        'block-device-mapping': { 'ami': 'vda' },
        'hostname': 'foo.bar.tld',
        'instance-action': 'none',
        'instance-id': '123456',
        'instance-type': 'h1.small',
        #'kernel-id': '',
        'local-hostname': 'foo.bar.tld',
        #'local-ipv4': '',
        'placement': {'availability-zone': 'bar'},
        'public-hostname': 'foo.bar.tld',
        #'public-ipv4': '',
        'public-keys': {
            'pk1': ['ssh-rsa ...'],
            'pk2': ['ssh-rsa ...']
        },
        'ramdisk-id': '',
        'reservation-id': '',
        'security-groups': ''
    }
}

endpoints = {
    '2009-04-04': meta,
    'latest': meta
}


@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def endepoints(path):
    """Catch all route for cloud-init EC2 DataSource that drills down the
    endpoints dictionary.
    """
    # split path and remove empty entries
    tokens = path.split('/')
    tokens = list(filter(None, tokens))

    # drill down the endpoints dictionary
    result = endpoints
    for item in tokens:
        """public keys need to be handled differently and returned in the
        form number=keyname (0=keyname) and respond to the following endpoints:
            - public-keys/
            - public-keys/0
            - public-keys/0/openssh-key
        """
        if (item == 'public-keys'):
            result = pks(result[item], tokens)
            break
        elif (item in result):
            result = result[item]

    """when the result is a string deliver it as the final value, otherwise
    return a dictionary representing the rest of the endpoints dict
    """
    if isinstance(result, str):
        output = result
    else:
        output = [r + '/' if isinstance(result[r], dict) else r for r in result.keys()]
        output = '\n'.join(output)

    return output


def pks(keys, tokens):
    # handle public keys special case and deliver them in the correct format
    sorted_keys = sorted(keys)

    if tokens[-1] == 'public-keys':
        output = {'{}={}'.format(idx, key) : '' for idx, key in enumerate(sorted_keys)}
    elif re.match('\d', tokens[-1]):
        output = 'openssh-key'
    elif tokens[-1] == 'openssh-key':
        key = sorted_keys[int(tokens[-2])]
        output = '\n'.join(keys[key])
    else:
        output = ''

    return output

if __name__ == '__main__':
    app.run(use_debugger=True, debug=True,
            use_reloader=True, host='0.0.0.0')