😱 Fix vulnerability of load_yaml and load_yamlf (#12)

owlmixin/util.py

@@ -9,7 +9,7 @@
 import json
 import yaml
 from urllib.request import urlopen
-from yaml import Loader, SafeLoader
+from yaml import SafeLoader
 
 import csv
 from csv import register_dialect, Dialect, QUOTE_MINIMAL
@@ -44,7 +44,6 @@ def increase_indent(self, flow=False, indentless=False):
 def construct_yaml_str(self, node):
     return self.construct_scalar(node)
 
-Loader.add_constructor(u'tag:yaml.org,2002:str', construct_yaml_str)
 SafeLoader.add_constructor(u'tag:yaml.org,2002:str', construct_yaml_str)
 
 
@@ -93,7 +92,7 @@ def load_yaml(yaml_str):
     :param unicode yaml_str:
     :rtype: dict | list
     """
-    return yaml.load(yaml_str)
+    return yaml.safe_load(yaml_str)
 
 
 def load_yamlf(fpath, encoding):
@@ -103,7 +102,7 @@ def load_yamlf(fpath, encoding):
     :rtype: dict | list
     """
     with codecs.open(fpath, encoding=encoding) as f:
-        return yaml.load(f)
+        return yaml.safe_load(f)
 
 
 def load_csvf(fpath, fieldnames, encoding):

tests/test_util.py

@@ -2,8 +2,11 @@
 
 from __future__ import division, absolute_import, unicode_literals
 
+from yaml.constructor import ConstructorError
 from owlmixin import util
 
+import pytest
+
 
 class TestReplaceKeys:
     def test_need_not_snake(self):
@@ -82,3 +85,38 @@ def test_docopt(self):
         assert util.to_snake("<file_list>") == "file_list"
         assert util.to_snake("-o") == "o"
         assert util.to_snake("--detail-option") == "detail_option"
+
+
+class TestLoadYaml:
+    def test(self):
+        actual = util.load_yaml('''
+id: 1
+names:
+  - tadashi
+  - aikawa
+        ''')
+        assert actual == {
+            "id": 1,
+            "names": ['tadashi', 'aikawa']
+        }
+
+    def test_yaml_load_vulnerability(self):
+        with pytest.raises(ConstructorError):
+            util.load_yaml('!!python/object/apply:os.system ["calc.exe"]')
+
+
+class TestLoadYamlf:
+    def test(self):
+        assert util.load_yamlf('tests/yaml/spots_utf8.yaml', 'utf-8') == [
+            {
+                "address": {"name": "address1"},
+                "names": ["spot1"]
+            },
+            {
+                "names": ["スポット21", "スポット22"]
+            }
+        ]
+
+    def test_yaml_load_vulnerability(self):
+        with pytest.raises(ConstructorError):
+            util.load_yamlf('tests/yaml/vulnerability.yaml', 'utf-8')

tests/yaml/vulnerability.yaml

@@ -0,0 +1 @@
+!!python/object/apply:os.system ["calc.exe"]