diff options
| -rw-r--r-- | cloudinit/safeyaml.py | 31 | ||||
| -rw-r--r-- | cloudinit/util.py | 3 | ||||
| -rw-r--r-- | tests/unittests/test_util.py | 39 |
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 |