cloud-init-cherrypy/main.py

import os
import sys

import cherrypy
from cherrypy.lib.static import serve_file
import yaml
import socket

import configparser

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 MainApp:
    def _init_ip(self):
        """
        Get remote IP
        """
        try:
            self.remoteip = cherrypy.request.headers.get(
                'X-Real-Ip',
                cherrypy.request.remote.ip
            )
        except BaseException:
            self.remoteip = cherrypy.request.remote.ip

        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")
        if os.path.exists(filepath):
            try:
                with open(filepath) as f:
                    content = f.read().splitlines()
                raise cherrypy.HTTPRedirect(content[0], 301)
            except BaseException:
                return False
        return False

    @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)
        return serve_file(filepath, "application/x-download", "attachment")

    @cherrypy.expose
    def meta_data(self):
        """
        Return meta-data in YAML
        """
        self._init_ip()
        self._redirect_if_needed()
        hostname = self.hostinfo[0]
        data = {
                "instance-id": hostname.split(".")[0],
                "local-hostname": hostname
        }

        filepath = os.path.join(PATH, "data", hostname, meta_data_filename)
        if os.path.exists(filepath):
            with open(filepath, "r") as metadata:
                for line in metadata.readlines():
                    linesplit = list(map(lambda k: k.strip(), line.split(":")))
                    data[linesplit[0]] = linesplit[1]

        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 f:
            f.write(data)


ROOT = MainApp()

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

    cherrypy.tree.mount(ROOT)
    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()