From 2d1eb2826ae8ba04d666173eb304ab0df7363dca Mon Sep 17 00:00:00 2001
From: Peter Surda <surda@economicsofbitcoin.com>
Date: Fri, 4 Jun 2021 18:54:17 +0800
Subject: [PATCH] Add support for #include user-data format

---
 main.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 59 insertions(+), 11 deletions(-)

diff --git a/main.py b/main.py
index 3437a54..b965ece 100644
--- a/main.py
+++ b/main.py
@@ -9,10 +9,10 @@ import socket
 import sys
 from ipaddress import AddressValueError, IPv4Address, IPv6Address
 
-import yaml
 import cherrypy
 from cherrypy.lib.static import serve_file
-
+from jinja2 import Template
+import yaml
 
 PATH = os.path.dirname(os.path.abspath(__file__))
 
@@ -74,6 +74,33 @@ class CloudInitApp:
                 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):
         """
@@ -85,7 +112,22 @@ class CloudInitApp:
                                 USER_DATA_FILENAME)
         if not os.path.exists(filepath):
             filepath = os.path.join(PATH, "data", USER_DATA_FILENAME)
-        return serve_file(filepath, "text/yaml", "attachment")
+        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):
@@ -94,13 +136,10 @@ class CloudInitApp:
         """
         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)
+        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:
                 for line in metadata.readlines():
@@ -110,7 +149,7 @@ class CloudInitApp:
         cherrypy.response.headers['Content-Type'] = \
             'text/yaml'
         cherrypy.response.headers['Content-Disposition'] = \
-            'attachment; filename="user-data"'
+            'attachment; filename="meta-data"'
         return yaml.dump(data)
 
     @cherrypy.expose
@@ -145,7 +184,16 @@ if __name__ == "__main__":
         CONFIG["server"].getint("server_port", 8081)
     ENGINE = cherrypy.engine
 
-    cherrypy.tree.mount(ROOT)
+    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"):