plugin recipe: improve recipe to correctly generate promise which has complex parameters

Use json.dumps/.loads to generate promise plugin code. This is safe as it will correctly escape sring and prevent code injection from untrustable parameters.

Add 'module' parameter to simplify parameters used to generate the script instead of passing full import code. When module is set, the generated script will import `module` from slapos.promise.plugin.

slapos/recipe/promise_plugin.py

@@ -24,19 +24,18 @@
 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 #
 ##############################################################################
-import logging, os, sys
+import json
+import re
+import logging, os
 import zc.buildout.easy_install
 from slapos.recipe.librecipe import GenericBaseRecipe
 
 script_template = '''# This script is auto generated by slapgrid, do not edit!
+import json
 import sys
-sys.path[0:0] = [
-  %(path)s
-]
+sys.path[0:0] = %(path)s
 
-extra_config_dict = {
-  %(config)s
-}
+extra_config_dict = json.loads("""%(config)s""")
 
 # We want to cleanup all imported modules from slapos namespace, because
 # they will conflict with slapos.core.
@@ -53,8 +52,6 @@ for module in sys.modules.keys():
   if 'slapos' in module or 'pkg_resources' in module:
     del sys.modules[module]
 
-import slapos.grid.promise
-
 %(content)s
 '''
 
@@ -79,11 +76,18 @@ class Recipe(GenericBaseRecipe):
     )
     if cache_storage is None:
       cache_storage = {}
+      try:
         setattr(
           self.buildout,
           self._WORKING_SET_CACHE_NAME,
           cache_storage
         )
+      except AttributeError:
+        if type(self.buildout) == type({}):
+          # failed to set attribute in test mode, cache not used
+          pass
+        else:
+          raise
     return cache_storage
 
   def install(self):
@@ -102,30 +106,43 @@ class Recipe(GenericBaseRecipe):
       develop_eggs_dir,
     )
     if cache_key not in cache_storage:
+      if develop_eggs_dir and eggs_dir:
         working_set = zc.buildout.easy_install.working_set(
           egg_list,
           [develop_eggs_dir, eggs_dir]
         )
         cache_storage[cache_key] = working_set
+      else:
+        working_set = set()
     else:
       working_set = cache_storage[cache_key]
 
-    content = self.options['content'].strip()
+    regex = r"^[\w_\-\.\s]+$"
+    module = self.options.get('module', '').strip()
+    if module:
+      if not re.search(regex, module):
+        raise ValueError("%r is not a valid module name" % module)
+      content_string = "from slapos.promise.plugin.%s import RunPromise" % module
+    else:
+      # old parameter for compatibility
+      content_string = self.options['content'].strip()
+      if not re.search(regex, module):
+        raise ValueError("Promise content %r is not valid" % content_string)
+
     output = self.options['output']
     mode = self.options.get('mode', '0600')
-    path_list_string = ""
+    path_list = []
     for dist in working_set:
-      path_list_string += '  "%s",\n' % dist.location
+      path_list.append(dist.location)
 
-    content_string = '\n'.join([line.lstrip() for line in content.split('\n')])
-    config_string = ""
+    config_dict = dict()
     for key in self.options:
       if key.startswith('config-'):
-        config_string += "  '%s': '%s',\n" % (key[7:], self.options[key])
+        config_dict[key[7:]] = self.options[key]
 
-    option_dict = dict(path=path_list_string.strip(),
+    option_dict = dict(path=json.dumps(path_list, indent=2),
                        content=content_string,
-                       config=config_string.strip())
+                       config=json.dumps(config_dict, indent=2, sort_keys=True))
     with open(output, 'w') as f:
       f.write(script_template % option_dict)
 

slapos/test/recipe/test_plugin.py

+import os, shutil, tempfile, unittest
+from slapos.recipe import promise_plugin
+from slapos.test.utils import makeRecipe
+import stat, json
+
+class TestPromisePlugin(unittest.TestCase):
+
+  def setUp(self):
+    self.tmp = tempfile.mkdtemp()
+    self.output = os.path.join(self.tmp, 'output.py')
+    self.options = options = {
+      'output': self.output,
+      'eggs': 'slapos.cookbook'
+    }
+
+  def tearDown(self):
+    shutil.rmtree(self.tmp)
+
+  def test_parameters(self):
+    self.options['mode'] = '0644'
+    self.options['module'] = 'check_site_available'
+    self.options['config-param1'] = "YY^@12"
+    self.options['config-param2'] = "23'91'"
+    self.options['config-param3'] = None
+    self.options['config-param4'] = """param
+in multi line
+123444
+"""
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    recipe.install()
+
+    self.assertTrue(os.path.exists(self.output))
+    with open(self.output) as f:
+      content = f.read()
+    self.assertIn("from slapos.promise.plugin.check_site_available import RunPromise", content)
+    self.assertEqual(stat.S_IMODE(os.stat(self.output).st_mode), int('644', 8))
+
+    expected_dict = dict(
+      param1=self.options['config-param1'],
+      param2=self.options['config-param2'],
+      param3=self.options['config-param3'],
+      param4=self.options['config-param4'],
+    )
+    self.assertIn('extra_config_dict = json.loads("""%s""")' % json.dumps(expected_dict, indent=2, sort_keys=True), content)
+
+  def test_no_module_set(self):
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    with self.assertRaises(KeyError):
+      recipe.install()
+
+  def test_default(self):
+    self.options['module'] = 'check_site_available'
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    recipe.install()
+
+    self.assertTrue(os.path.exists(self.output))
+    self.assertEqual(stat.S_IMODE(os.stat(self.output).st_mode), int('600', 8))
+    with open(self.output) as f:
+      content = f.read()
+    self.assertIn("from slapos.promise.plugin.check_site_available import RunPromise", content)
+    self.assertIn('extra_config_dict = json.loads("""{}""")', content)
+
+
+  def test_bad_parameters(self):
+    self.options['module'] = 'check_site_available'
+    self.options['config-param1; print "toto"'] = """#xxxx"\nimport os; os.stat(f)"""
+    self.options['config-param2\n@domething'] = '"#$$*PPP\n\n p = 2*5; print "result is %s" % p'
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    recipe.install()
+
+    self.assertTrue(os.path.exists(self.output))
+    with open(self.output) as f:
+      content = f.read()
+
+    expected_param1 = '"param1; print \\"toto\\"": "#xxxx\\"\\nimport os; os.stat(f)",'
+    expected_param2 = '"param2\\n@domething": "\\"#$$*PPP\\n\\n p = 2*5; print \\"result is %s\\" % p"'
+    self.assertIn(expected_param1, content)
+    self.assertIn(expected_param2, content)
+
+  def test_bad_module_name(self):
+    self.options['module'] = 'check_site_available; print "toto"'
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    with self.assertRaises(ValueError) as p:
+      recipe.install()
+
+    self.assertEqual(p.exception.message, "%r is not a valid module name" % self.options['module'])
+
+  def test_bad_content(self):
+    self.options['content'] = 'from slapos.plugin.check_site_available import toto; print "toto"'
+    recipe = makeRecipe(
+            promise_plugin.Recipe,
+            options=self.options,
+            name='plugin')
+    with self.assertRaises(ValueError) as p:
+      recipe.install()
+
+    self.assertEqual(p.exception.message, "Promise content %r is not valid" % self.options['content'])
+