Self parse user data

stacker/blueprints/base.py

@@ -1,11 +1,13 @@
 import copy
 import hashlib
 import logging
+import re
+import base64
 
 from troposphere import (
     Parameter,
     Ref,
-    Template,
+    Template
 )
 
 from ..exceptions import (
@@ -221,6 +223,7 @@ def resolve_variable(var_name, var_def, provided_variable, blueprint_name):
 
 
 class Blueprint(object):
+
     """Base implementation for rendering a troposphere template.
 
     Args:
@@ -402,6 +405,46 @@ def render_template(self):
         version = hashlib.md5(rendered).hexdigest()[:8]
         return (version, rendered)
 
+    def parse_user_data(self, raw_user_data):
+        """Translate a userdata file to into the file contents.
+
+        It supports referencing template variables to create userdata
+        that's supplemented with information from the data, as commonly
+        required when creating EC2 userdata files. Automatically, encodes
+        the data file to base64 after it is processed.
+
+        Args:
+            raw_user_data (str): the user data with the cloud-init info
+
+        Returns:
+            str: The parsed user data, with all the variables values and
+                 refs replaced with their resolved values.
+
+        """
+        pattern = re.compile(r'{{([::|\w]+)}}')
+        res = ""
+        start_index = 0
+        variables = self.get_variables()
+
+        for match in pattern.finditer(raw_user_data):
+            res += raw_user_data[start_index:match.start()]
+
+            key = match.group(1)
+
+            if key in variables:
+                if type(variables[key]) is CFNParameter:
+                    res += variables[key].to_parameter_value()
+                else:
+                    res += variables[key]
+            else:
+                raise MissingVariable(self.name, key)
+
+            start_index = match.end()
+
+        res += raw_user_data[start_index:]
+
+        return base64.b64encode(res)
+
     @property
     def rendered(self):
         if not self._rendered:

stacker/tests/blueprints/test_base.py

@@ -1,10 +1,11 @@
 import unittest
+import base64
 
 from mock import MagicMock
 from troposphere import (
     Base64,
     Ref,
-    s3,
+    s3
 )
 
 from stacker.blueprints.base import (
@@ -44,6 +45,7 @@ def mock_lookup_handler(value, provider=None, context=None, fqn=False,
 
 
 class TestBuildParameter(unittest.TestCase):
+
     def test_base_parameter(self):
         p = build_parameter("BasicParam", {"type": "String"})
         p.validate()
@@ -72,6 +74,7 @@ class TestBlueprint(Blueprint):
             }
 
         class TestBlueprintSublcass(TestBlueprint):
+
             def defined_variables(self):
                 variables = super(TestBlueprintSublcass,
                                   self).defined_variables()
@@ -528,3 +531,44 @@ class TestBlueprint(Blueprint):
 
         with self.assertRaises(AttributeError):
             TestBlueprint(name="test", context=MagicMock())
+
+    def test_parse_user_data(self):
+        class TestBlueprint(Blueprint):
+            VARIABLES = {
+                "Param1": {"type": str}
+            }
+
+        blueprint = TestBlueprint(name="test", context=MagicMock())
+        variables = [Variable("Param1", "test1")]
+        blueprint.resolve_variables(variables)
+        userdata_raw = "Param1: {{Param1}}"
+        userdata = blueprint.parse_user_data(userdata_raw)
+        print(userdata)
+        self.assertEqual(userdata, base64.b64encode("Param1: test1"))
+
+    def test_parse_user_data_cfn_parameters(self):
+        class TestBlueprint(Blueprint):
+            VARIABLES = {
+                "Param2": {"type": CFNString},
+            }
+
+        blueprint = TestBlueprint(name="test", context=MagicMock())
+        variables = [Variable("Param1", "test1"), Variable("Param2", "test2")]
+        blueprint.resolve_variables(variables)
+        userdata_raw = "Param2: {{Param2}}!"
+        userdata = blueprint.parse_user_data(userdata_raw)
+        self.assertEqual(userdata, base64.b64encode("Param2: test2!"))
+
+    def test_parse_user_data_fails(self):
+        class TestBlueprint(Blueprint):
+            VARIABLES = {
+                "Param1": {"type": str},
+                "Param2": {"type": str},
+            }
+
+        blueprint = TestBlueprint(name="test", context=MagicMock())
+        variables = [Variable("Param1", "test1"), Variable("Param2", "test2")]
+        blueprint.resolve_variables(variables)
+        userdata_raw = "My name is {{Param1}} and {{Param3}}."
+        with self.assertRaises(MissingVariable):
+            blueprint.parse_user_data(userdata_raw)