From bc795757febdcce430d89f9d08f75c32d6989d3c Mon Sep 17 00:00:00 2001 From: Thane Thomson Date: Tue, 7 Nov 2017 13:33:54 +0200 Subject: [PATCH] attempting to fix security flaw (issue #1) --- mlalchemy/parser.py | 2 +- tests/test_yaml_security.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) create mode 100644 tests/test_yaml_security.py diff --git a/mlalchemy/parser.py b/mlalchemy/parser.py index 44f2711..faeecdd 100644 --- a/mlalchemy/parser.py +++ b/mlalchemy/parser.py @@ -33,7 +33,7 @@ def parse_yaml_query(yaml_content): On success, the processed MLQuery object. """ logger.debug("Attempting to parse YAML content:\n%s" % yaml_content) - return parse_query(yaml.load(yaml_content)) + return parse_query(yaml.safe_load(yaml_content)) def parse_json_query(json_content): diff --git a/tests/test_yaml_security.py b/tests/test_yaml_security.py new file mode 100644 index 0000000..18d9337 --- /dev/null +++ b/tests/test_yaml_security.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +from __future__ import unicode_literals + +import unittest +import yaml + +from mlalchemy import * +from mlalchemy.testing import MLAlchemyTestCase + + +class TestYamlSecurity(MLAlchemyTestCase): + + def test_basic_yaml_security(self): + with self.assertRaises(yaml.constructor.ConstructorError): + parse_yaml_query('!!python/object/apply:os.system ["echo Hello"]') + + +if __name__ == "__main__": + unittest.main() +