cloud-init-cherrypy/main.py

#!/usr/bin/env python3
"""
Serve cloud init files
"""

import configparser
import os
import socket
import sys
from ipaddress import AddressValueError, IPv4Address, IPv6Address

import cherrypy
from cherrypy.lib.static import serve_file
from jinja2 import Template
import yaml

PATH = os.path.dirname(os.path.abspath(__file__))

CONFIG = configparser.ConfigParser()
CONFIG.read(os.path.join(PATH, "config.ini"))

USER_DATA_FILENAME = CONFIG["app"].get("user_data", "user-data")
META_DATA_FILENAME = CONFIG["app"].get("meta_data", "meta-data")
REDIRECT_FILENAME = CONFIG["app"].get("redirect", "redirect")


class CloudInitApp:
    """
    Serve cloud init files
    """

    def __init__(self):
        self.remoteip = None
        self.hostinfo = ('localhost', )

    def _can_ip_be_proxy(self):
        self.remoteip = cherrypy.request.remote.ip
        try:
            ipobj = IPv4Address(self.remoteip)
        except AddressValueError:
            try:
                ipobj = IPv6Address(self.remoteip)
            except AddressValueError:
                return False
        return not ipobj.is_global

    def _init_ip(self):
        """
        Get remote IP
        """
        if self._can_ip_be_proxy():
            try:
                self.remoteip = cherrypy.request.headers.get(
                    'X-Real-Ip',
                    cherrypy.request.remote.ip
                )
            except KeyError:
                pass

        try:
            self.hostinfo = socket.gethostbyaddr(self.remoteip)
        except socket.herror:
            self.hostinfo = ('localhost', )

    def _redirect_if_needed(self):
        filepath = os.path.join(PATH, "data", self.hostinfo[0],
                                REDIRECT_FILENAME)
        if os.path.exists(filepath):
            try:
                with open(filepath) as redirect:
                    content = redirect.read().splitlines()
                raise cherrypy.HTTPRedirect(content[0], 301)
            except IOError:
                return False
        return False

    def _generate_metadata(self):
        self._init_ip()
        hostname = self.hostinfo[0]
        base = cherrypy.request.base
        data = {
            "instance-id": hostname.split(".")[0],
            "local-hostname": hostname,
            "baseurl": base
        }
        return data

    @staticmethod
    def _content_type(data):
        if data.startswith("#include"):
            return "text/x-include-url"
        elif data.startswith("## template: jinja"):
            return "text/jinja2"
        else:
            return "text/cloud-config"

    @staticmethod
    def _wrap_metadata(metadata):
        return {'ds': {
                          'meta_data': metadata
                         }
                  }

    @cherrypy.expose
    def user_data(self):
        """
        Serves a static file
        """
        self._init_ip()
        self._redirect_if_needed()
        filepath = os.path.join(PATH, "data", self.hostinfo[0],
                                USER_DATA_FILENAME)
        if not os.path.exists(filepath):
            filepath = os.path.join(PATH, "data", USER_DATA_FILENAME)
        with open(filepath, "r") as userdata:
            data = userdata.read()
 
        c = CloudInitApp._content_type(data)

        cherrypy.response.headers['Content-Type'] = c
        cherrypy.response.headers['Content-Disposition'] = \
            'attachment; filename="user-data"'

        if c == "text/x-include-url":
            t = Template(data)
            metadata = self._generate_metadata()
            wrapped = CloudInitApp._wrap_metadata(metadata)
            return t.render(**wrapped)
        else:
            return data

    @cherrypy.expose
    def meta_data(self):
        """
        Return meta-data in YAML
        """
        self._init_ip()
        self._redirect_if_needed()

        data = self._generate_metadata()

        filepath = os.path.join(PATH, "data", data['local-hostname'], META_DATA_FILENAME)
        if os.path.exists(filepath):
            with open(filepath, "r") as metadata:
                data.update(yaml.safe_load(metadata))

        cherrypy.response.headers['Content-Type'] = \
            'text/yaml'
        cherrypy.response.headers['Content-Disposition'] = \
            'attachment; filename="meta-data"'
        return yaml.dump(data)

    @cherrypy.expose
    def vendor_data(self):
        """
        Return empty vendor-data
        """
        return ""

    @cherrypy.expose
    def finished(self, data):
        """
        Saves additional meta-data

        :param data: meta-data to be added
        """
        self._init_ip()
        folder = os.path.join(PATH, "data", self.hostinfo[0])
        if not os.path.exists(folder):
            os.makedirs(folder)

        with open(os.path.join(folder, META_DATA_FILENAME), "w") as fin:
            fin.write(data)


ROOT = CloudInitApp()

if __name__ == "__main__":
    cherrypy.server.socket_host = \
        CONFIG["server"].get("server_host", "127.0.0.1")
    cherrypy.server.socket_port = \
        CONFIG["server"].getint("server_port", 8081)
    ENGINE = cherrypy.engine

    current_dir = os.path.dirname(os.path.abspath(__file__))
    config = {
              '/include': {
                           'tools.staticdir.on': True,
                           'tools.staticdir.dir': os.path.join(current_dir, 'data', 'include'),
                           'tools.staticdir.content_types': {'yml': 'text/yaml'}
                          }
             }
    cherrypy.tree.mount(ROOT, config=config)

    if hasattr(ENGINE, "signal_handler"):
        ENGINE.signal_handler.subscribe()
    if hasattr(ENGINE, "console_control_handler"):
        ENGINE.console_control_handler.subscribe()
    try:
        ENGINE.start()
    except Exception:
        sys.exit(1)
    else:
        ENGINE.block()