attempting to fix security flaw (issue #1)

Author: thanethomson

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):

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()
+