AMBARI-18366 - YAML Maps For Storm Are Not Being Escaped Correctly (jonathanhurley)

ambari-common/src/main/python/ambari_commons/yaml_utils.py

@@ -19,37 +19,60 @@ limitations under the License.
 """
 import re
 
+# [a,b,c]
+REGEX_LIST = '^\w*\[.+\]\w*$'
+
+# {a: v, b: v2, c: v3}
+REGEX_DICTIONARY = '^\w*\{.+\}\w*$'
+
+"""
+storm:
+  hosts:
+    [c6401.ambari.apache.org, c6402.ambari.apache.org]
+  groups:
+    [hadoop, foo]
+"""
+REGEX_NESTED_MAPS = '^[\w+\s*:\s*\n\s*]+\[(.*?)\]+'
+
 def escape_yaml_property(value):
-  unquouted = False
   unquouted_values = ["null", "Null", "NULL", "true", "True", "TRUE", "false",
     "False", "FALSE", "YES", "Yes", "yes", "NO", "No", "no", "ON", "On", "on",
     "OFF", "Off", "off"]
 
+  # known list of boolean/null types
   if value in unquouted_values:
-    unquouted = True
-
-  # if is list [a,b,c] or dictionary {a: v, b: v2, c: v3}
-  if re.match('^\w*\[.+\]\w*$', value) or re.match('^\w*\{.+\}\w*$', value):
-    unquouted = True
+    return value
 
+  # quick pythonic check for integer
   try:
     int(value)
-    unquouted = True
+    return value
   except ValueError:
     pass
 
+  # quick pythonic check for float
   try:
     float(value)
-    unquouted = True
+    return value
   except ValueError:
     pass
 
-  if not unquouted:
-    value = value.replace("'", "''")
-    value = "'" + value + "'"
+  # if is list [a,b,c] or dictionary {a: v, b: v2, c: v3}
+  if re.match(REGEX_LIST, value) or re.match(REGEX_DICTIONARY, value):
+    return value
 
+  # check for a nested map
+  if re.match(REGEX_NESTED_MAPS, value):
+    # nested maps must begin on a newline and not have whitespace on the first line
+    value = value.lstrip()
+    return "\n" + value
+
+  # no more checks, so assume it's a string a quote it
+  value = value.replace("'", "''")
+  value = "'" + value + "'"
   return value
 
+
 def get_values_from_yaml_array(yaml_array):
   """
   Converts a YAML array into a normal array of values. For example, this

ambari-server/src/test/python/TestYAMLUtils.py

@@ -41,4 +41,42 @@ class TestYAMLUtils(TestCase):
     self.assertEquals(expected_values, values)
 
     values = yaml_utils.get_values_from_yaml_array('[\'c6401.ambari.apache.org\', "c6402.ambari.apache.org"]')
-    self.assertEquals(expected_values, values)
+    self.assertEquals(expected_values, values)
+
+
+  def test_yaml_property_escaping(self):
+    """
+    Tests that YAML values are escaped with quotes properly when needed
+    """
+    self.assertEquals("True", yaml_utils.escape_yaml_property("True"))
+    self.assertEquals("FALSE", yaml_utils.escape_yaml_property("FALSE"))
+    self.assertEquals("yes", yaml_utils.escape_yaml_property("yes"))
+    self.assertEquals("NO", yaml_utils.escape_yaml_property("NO"))
+    self.assertEquals("28", yaml_utils.escape_yaml_property("28"))
+    self.assertEquals("28.0", yaml_utils.escape_yaml_property("28.0"))
+    self.assertEquals("[a,b,c]", yaml_utils.escape_yaml_property("[a,b,c]"))
+    self.assertEquals("{ foo : bar }", yaml_utils.escape_yaml_property("{ foo : bar }"))
+
+    # some strings which should be escaped
+    self.assertEquals("'5f'", yaml_utils.escape_yaml_property("5f"))
+    self.assertEquals("'28.O'", yaml_utils.escape_yaml_property("28.O"))
+    self.assertEquals("'This is a test of a string'", yaml_utils.escape_yaml_property("This is a test of a string"))
+
+    # test maps
+    map = """
+      storm:
+        hosts:
+          [c6401.ambari.apache.org, c6402.ambari.apache.org]
+        groups:
+          [hadoop, foo]
+        foo:
+          [bar, baz]
+        foo2:
+          bar2:
+            [baz2]
+    """
+    escaped_map = yaml_utils.escape_yaml_property(map)
+    self.assertTrue(escaped_map.startswith("\n"))
+    self.assertFalse("'" in escaped_map)
+
+