add 'safeyaml' to cloudinit

In 0.7.0 we started using yaml.safe_load to load data rather than
yaml.load.  Some producers (namely, ubuntu MAAS created) have produced
cloud-config data in the past that included python unicode types.

This creates a specialized safe_loader that is basically safe_load +
support for python unicode.

3 files changed, 72 insertions, 1 deletions

diff --git a/cloudinit/safeyaml.py b/cloudinit/safeyaml.py
new file mode 100644
index 00000000..8b4da1fa
--- /dev/null
+++ b/cloudinit/safeyaml.py
@@ -0,0 +1,31 @@
+# vi: ts=4 expandtab
+#
+#    Copyright (C) 2012 Canonical Ltd.
+#
+#    Author: Scott Moser <scott.moser@canonical.com>
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License version 3, as
+#    published by the Free Software Foundation.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import yaml
+
+
+class _CustomSafeLoader(yaml.SafeLoader):
+    def construct_python_unicode(self, node):
+        return self.construct_scalar(node)
+
+_CustomSafeLoader.add_constructor(
+    u'tag:yaml.org,2002:python/unicode',
+    _CustomSafeLoader.construct_python_unicode)
+
+def load(blob):
+    return(yaml.load(blob, Loader=_CustomSafeLoader))
diff --git a/cloudinit/util.py b/cloudinit/util.py
index 94b17dfa..46d490f7 100644
--- a/cloudinit/util.py
+++ b/cloudinit/util.py
@@ -50,6 +50,7 @@ import yaml
 
 from cloudinit import importer
 from cloudinit import log as logging
+from cloudinit import safeyaml
 from cloudinit import url_helper as uhelp
 
 from cloudinit.settings import (CFG_BUILTIN)
@@ -612,7 +613,7 @@ def load_yaml(blob, default=None, allowed=(dict,)):
         LOG.debug(("Attempting to load yaml from string "
                  "of length %s with allowed root types %s"),
                  len(blob), allowed)
-        converted = yaml.safe_load(blob)
+        converted = safeyaml.load(blob)
         if not isinstance(converted, allowed):
             # Yes this will just be caught, but thats ok for now...
             raise TypeError(("Yaml load allows %s root types,"
diff --git a/tests/unittests/test_util.py b/tests/unittests/test_util.py
index 15fcbd26..96962b91 100644
--- a/tests/unittests/test_util.py
+++ b/tests/unittests/test_util.py
@@ -1,5 +1,6 @@
 import os
 import stat
+import yaml
 
 from mocker import MockerTestCase
 from unittest import TestCase
@@ -268,4 +269,42 @@ class TestGetCmdline(TestCase):
         os.environ['DEBUG_PROC_CMDLINE'] = 'abcd 123'
         self.assertEqual(os.environ['DEBUG_PROC_CMDLINE'], util.get_cmdline())
 
+
+class TestLoadYaml(TestCase):
+    mydefault = "7b03a8ebace993d806255121073fed52"
+
+    def test_simple(self):
+        mydata = {'1': "one", '2': "two"}
+        self.assertEqual(util.load_yaml(yaml.dump(mydata)), mydata)
+
+    def test_nonallowed_returns_default(self):
+        # for now, anything not in the allowed list just returns the default.
+        myyaml = yaml.dump({'1': "one"})
+        self.assertEqual(util.load_yaml(blob=myyaml,
+                                        default=self.mydefault,
+                                        allowed=(str,)),
+                         self.mydefault)
+
+    def test_bogus_returns_default(self):
+        badyaml = "1\n 2:"
+        self.assertEqual(util.load_yaml(blob=badyaml,
+                                        default=self.mydefault),
+                         self.mydefault)
+
+    def test_unsafe_types(self):
+        # should not load complex types
+        unsafe_yaml = yaml.dump((1, 2, 3,))
+        self.assertEqual(util.load_yaml(blob=unsafe_yaml,
+                                        default=self.mydefault),
+                         self.mydefault)
+
+    def test_python_unicode(self):
+        # complex type of python/unicde is explicitly allowed
+        myobj = {'1': unicode("FOOBAR")}
+        safe_yaml = yaml.dump(myobj)
+        self.assertEqual(util.load_yaml(blob=safe_yaml,
+                                        default=self.mydefault),
+                         myobj)
+
+
 # vi: ts=4 expandtab