#!/usr/bin/env python3

# Provides a somewhat random, somewhat compat, somewhat useful mock version of

# http://docs.amazonwebservices.com

#   /AWSEC2/2007-08-29/DeveloperGuide/AESDG-chapter-instancedata.htm

"""

To use this to mimic the EC2 metadata service entirely, run it like:

  # Where 'eth0' is *some* interface.

  sudo ifconfig eth0:0 169.254.169.254 netmask 255.255.255.255

  sudo ./mock-meta.py -a 169.254.169.254 -p 80

Then:

  wget -q http://169.254.169.254/latest/meta-data/instance-id -O -; echo

  curl --silent http://169.254.169.254/latest/meta-data/instance-id ; echo

  ec2metadata --instance-id

"""

import argparse

import functools

import json

import logging

import os

import random

import socket

import string

import sys

import yaml

try:

    from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler

    import httplib as hclient

except ImportError:

    from http.server import HTTPServer, BaseHTTPRequestHandler

    from http import client as hclient

log = logging.getLogger('meta-server')

EC2_VERSIONS = [

    '1.0',

    '2007-01-19',

    '2007-03-01',

    '2007-08-29',

    '2007-10-10',

    '2007-12-15',

    '2008-02-01',

    '2008-09-01',

    '2009-04-04',

]

BLOCK_DEVS = [

    'ami',

    'ephemeral0',

    'root',

]

DEV_PREFIX = 'v'  # This seems to vary alot depending on images...

DEV_MAPPINGS = {

    'ephemeral0': '%sda2' % (DEV_PREFIX),

    'root': '/dev/%sda1' % (DEV_PREFIX),

    'ami': '%sda1' % (DEV_PREFIX),

    'swap': '%sda3' % (DEV_PREFIX),

}

META_CAPABILITIES = [

    'aki-id',

    'ami-id',

    'ami-launch-index',

    'ami-manifest-path',

    'ari-id',

    'block-device-mapping/',

    'hostname',

    'instance-action',

    'instance-id',

    'instance-type',

    'local-hostname',

    'local-ipv4',

    'placement/',

    'product-codes',

    'public-hostname',

    'public-ipv4',

    'public-keys/',

    'reservation-id',

    'security-groups'

]

PUB_KEYS = {

    'brickies': [

        ('ssh-rsa '

         'AAAAB3NzaC1yc2EAAAABIwAAAQEA3I7VUf2l5gSn5uavROsc5HRDpZdQueUq5ozemN'

         'Sj8T7enqKHOEaFoU2VoPgGEWC9RyzSQVeyD6s7APMcE82EtmW4skVEgEGSbDc1pvxz'

         'xtchBj78hJP6Cf5TCMFSXw+Fz5rF1dR23QDbN1mkHs7adr8GW4kSWqU7Q7NDwfIrJJ'

         'tO7Hi42GyXtvEONHbiRPOe8stqUly7MvUoN+5kfjBM8Qqpfl2+FNhTYWpMfYdPUnE7'

         'u536WqzFmsaqJctz3gBxH9Ex7dFtrxR4qiqEr9Qtlu3xGn7Bw07/+i1D+ey3ONkZLN'

         '+LQ714cgj8fRS4Hj29SCmXp5Kt5/82cD/VN3NtHw== brickies'),

        '',

    ],

}

INSTANCE_TYPES = [

    'm1.large',

    'm1.medium',

    'm1.small',

    'm1.xlarge',

]

AVAILABILITY_ZONES = [

    "us-east-1a",

    "us-east-1b",

    "us-east-1c",

    "us-east-1d",

    'eu-west-1a',

    'eu-west-1b',

    'us-west-1',

]

PLACEMENT_CAPABILITIES = {

    'availability-zone': AVAILABILITY_ZONES,

}

NOT_IMPL_RESPONSE = json.dumps({})

class WebException(Exception):

    def __init__(self, code, msg):

        Exception.__init__(self, msg)

        self.code = code

def yamlify(data):

    formatted = yaml.dump(data,

                          line_break="\n",

                          indent=4,

                          explicit_start=True,

                          explicit_end=True,

                          default_flow_style=False)

    return formatted

def format_text(text):

    if not len(text):

        return "<<"

    lines = text.splitlines()

    nlines = []

    for line in lines:

        nlines.append("<< %s" % line)

    return "\n".join(nlines)

def traverse(keys, mp):

    result = dict(mp)

    for k in keys:

        try:

            result = result.get(k)

        except (AttributeError, TypeError):

            result = None

            break

    return result

ID_CHARS = [c for c in (string.ascii_uppercase + string.digits)]

def id_generator(size=6, lower=False):

    txt = ''.join(random.choice(ID_CHARS) for x in range(size))

    if lower:

        return txt.lower()

    else:

        return txt

def get_ssh_keys():

    keys = {}

    keys.update(PUB_KEYS)

    # Nice helper to add in the 'running' users key (if they have one)

    key_pth = os.path.expanduser('~/.ssh/id_rsa.pub')

    if not os.path.isfile(key_pth):

        key_pth = os.path.expanduser('~/.ssh/id_dsa.pub')

    if os.path.isfile(key_pth):

        with open(key_pth, 'rb') as fh:

            contents = fh.read()

        keys[os.getlogin()] = [contents, '']

    return keys

class HTTPServerV6(HTTPServer):

    address_family = socket.AF_INET6

class MetaDataHandler(object):

    def __init__(self, opts):

        self.opts = opts

        self.instances = {}

    def get_data(self, params, who, **kwargs):

        if not params:

            # Show the root level capabilities when

            # no params are passed...

            caps = sorted(META_CAPABILITIES)

            return "\n".join(caps)

        action = params[0]

        action = action.lower()

        if action == 'instance-id':

            return 'i-%s' % (id_generator(lower=True))

        elif action == 'ami-launch-index':

            return "%s" % random.choice([0, 1, 2, 3])

        elif action == 'aki-id':

            return 'aki-%s' % (id_generator(lower=True))

        elif action == 'ami-id':

            return 'ami-%s' % (id_generator(lower=True))

        elif action == 'ari-id':

            return 'ari-%s' % (id_generator(lower=True))

        elif action == 'block-device-mapping':

            nparams = params[1:]

            if not nparams:

                return "\n".join(BLOCK_DEVS)

            else:

                subvalue = traverse(nparams, DEV_MAPPINGS)

                if not subvalue:

                    return "\n".join(sorted(list(DEV_MAPPINGS.keys())))

                else:

                    return str(subvalue)

        elif action in ['hostname', 'local-hostname', 'public-hostname']:

            # Just echo back there own hostname that they called in on..

            return "%s" % (who)

        elif action == 'instance-type':

            return random.choice(INSTANCE_TYPES)

        elif action == 'ami-manifest-path':

            return 'my-amis/spamd-image.manifest.xml'

        elif action == 'security-groups':

            return 'default'

        elif action in ['local-ipv4', 'public-ipv4']:

            # Just echo back there own ip that they called in on...

            return "%s" % (kwargs.get('client_ip', '10.0.0.1'))

        elif action == 'reservation-id':

            return "r-%s" % (id_generator(lower=True))

        elif action == 'product-codes':

            return "%s" % (id_generator(size=8))

        elif action == 'public-keys':

            nparams = params[1:]

            # This is a weird kludge, why amazon why!!!

            # public-keys is messed up, list of /latest/meta-data/public-keys/

            # shows something like: '0=brickies'

            # but a GET to /latest/meta-data/public-keys/0=brickies will fail

            # you have to know to get '/latest/meta-data/public-keys/0', then

            # from there you get a 'openssh-key', which you can get.

            # this hunk of code just re-works the object for that.

            avail_keys = get_ssh_keys()

            key_ids = sorted(list(avail_keys.keys()))

            if nparams:

                mybe_key = nparams[0]

                try:

                    key_id = int(mybe_key)

                    key_name = key_ids[key_id]

                except ValueError as e:

                    raise WebException(

                        hclient.BAD_REQUEST, "%s: not an integer" % mybe_key

                    ) from e

                except IndexError as e:

                    raise WebException(

                        hclient.NOT_FOUND, "Unknown key id %r" % mybe_key

                    ) from e

                # Extract the possible sub-params

                result = traverse(nparams[1:], {

                    "openssh-key": "\n".join(avail_keys[key_name]),

                })

                if isinstance(result, (dict)):

                    # TODO(harlowja): This might not be right??

                    result = "\n".join(sorted(result.keys()))

                if not result:

                    result = ''

                return result

            else:

                contents = []

                for (i, key_id) in enumerate(key_ids):

                    contents.append("%s=%s" % (i, key_id))

                return "\n".join(contents)

        elif action == 'placement':

            nparams = params[1:]

            if not nparams:

                pcaps = sorted(PLACEMENT_CAPABILITIES.keys())

                return "\n".join(pcaps)

            else:

                pentry = nparams[0].strip().lower()

                if pentry == 'availability-zone':

                    zones = PLACEMENT_CAPABILITIES[pentry]

                    return "%s" % random.choice(zones)

                else:

                    return "%s" % (PLACEMENT_CAPABILITIES.get(pentry, ''))

        else:

            log.warning(("Did not implement action %s, "

                         "returning empty response: %r"),

                        action, NOT_IMPL_RESPONSE)

            return NOT_IMPL_RESPONSE

class UserDataHandler(object):

    def __init__(self, opts):

        self.opts = opts

    def _get_user_blob(self, **kwargs):

        blob = None

        if self.opts['user_data_file'] is not None:

            blob = self.opts['user_data_file']

        if not blob:

            blob_mp = {

                'hostname': kwargs.get('who', 'localhost'),

            }

            lines = [

                "#cloud-config",

                yamlify(blob_mp),

            ]

            blob = "\n".join(lines)

        return blob.strip()

    def get_data(self, params, who, **kwargs):

        if not params:

            return self._get_user_blob(who=who)

        return NOT_IMPL_RESPONSE

# Seem to need to use globals since can't pass

# data into the request handlers instances...

# Puke!

meta_fetcher = None

user_fetcher = None

class Ec2Handler(BaseHTTPRequestHandler):

    def _get_versions(self):

        versions = ['latest'] + EC2_VERSIONS

        versions = sorted(versions)

        return "\n".join(versions)

    def log_message(self, fmt, *args):

        msg = "%s - %s" % (self.address_string(), fmt % (args))

        log.info(msg)

    def _find_method(self, path):

        # Puke! (globals)

        func_mapping = {

            'user-data': user_fetcher.get_data,

            'meta-data': meta_fetcher.get_data,

        }

        segments = [piece for piece in path.split('/') if len(piece)]

        log.info("Received segments %s", segments)

        if not segments:

            return self._get_versions

        date = segments[0].strip().lower()

        if date not in self._get_versions():

            raise WebException(hclient.BAD_REQUEST,

                               "Unknown version format %r" % date)

        if len(segments) < 2:

            raise WebException(hclient.BAD_REQUEST, "No action provided")

        look_name = segments[1].lower()

        if look_name not in func_mapping:

            raise WebException(hclient.BAD_REQUEST,

                               "Unknown requested data %r" % look_name)

        base_func = func_mapping[look_name]

        who = self.address_string()

        ip_from = self.client_address[0]

        if who == ip_from:

            # Nothing resolved, so just use 'localhost'

            who = 'localhost'

        kwargs = {

            'params': list(segments[2:]),

            'who': who,

            'client_ip': ip_from,

        }

        return functools.partial(base_func, **kwargs)

    def _do_response(self):

        who = self.client_address

        log.info("Got a call from %s for path %s", who, self.path)

        try:

            func = self._find_method(self.path)

            data = func()

            if not data:

                data = ''

            self.send_response(hclient.OK)

            self.send_header("Content-Type", "binary/octet-stream")

            self.send_header("Content-Length", len(data))

            log.info("Sending data (len=%s):\n%s", len(data),

                     format_text(data))

            self.end_headers()

            self.wfile.write(data.encode())

        except RuntimeError as e:

            log.exception("Error somewhere in the server.")

            self.send_error(hclient.INTERNAL_SERVER_ERROR, message=str(e))

        except WebException as e:

            code = e.code

            log.exception(str(e))

            self.send_error(code, message=str(e))

    def do_GET(self):

        self._do_response()

    def do_POST(self):

        self._do_response()

def setup_logging(log_level, fmt='%(levelname)s: @%(name)s : %(message)s'):

    root_logger = logging.getLogger()

    console_logger = logging.StreamHandler(sys.stdout)

    console_logger.setFormatter(logging.Formatter(fmt))

    root_logger.addHandler(console_logger)

    root_logger.setLevel(log_level)

def extract_opts():

    parser = argparse.ArgumentParser()

    parser.add_argument("-p", "--port", dest="port", action="store", type=int,

                        default=80, metavar="PORT",

                        help=("port from which to serve traffic"

                              " (default: %default)"))

    parser.add_argument("-a", "--addr", dest="address", action="store",

                        type=str, default='::', metavar="ADDRESS",

                        help=("address from which to serve traffic"

                              " (default: %default)"))

    parser.add_argument("-f", '--user-data-file', dest='user_data_file',

                        action='store', metavar='FILE',

                        help=("user data filename to serve back to"

                              "incoming requests"))

    parser.add_argument('extra', nargs='*')

    args = parser.parse_args()

    out = {'port': args.port, 'address': args.address, 'extra': args.extra,

           'user_data_file': None}

    if args.user_data_file:

        if not os.path.isfile(args.user_data_file):

            parser.error("Option -f specified a non-existent file")

        with open(args.user_data_file, 'rb') as fh:

            out['user_data_file'] = fh.read()

    return out

def setup_fetchers(opts):

    global meta_fetcher

    global user_fetcher

    meta_fetcher = MetaDataHandler(opts)

    user_fetcher = UserDataHandler(opts)

def run_server():

    # Using global here since it doesn't seem like we

    # can pass opts into a request handler constructor...

    opts = extract_opts()

    setup_logging(logging.DEBUG)

    setup_fetchers(opts)

    log.info("CLI opts: %s", opts)

    server_address = (opts['address'], opts['port'])

    server = HTTPServerV6(server_address, Ec2Handler)

    sa = server.socket.getsockname()

    log.info("Serving ec2 metadata on %s using port %s ...", sa[0], sa[1])

    server.serve_forever()

if __name__ == '__main__':

    run_server()

# vi: ts=4 expandtab