WIP: Another safe_load API attempt

examples/safe.py

@@ -0,0 +1,30 @@
+import yaml
+import os
+
+def myconstructor(loader, node):
+    value = loader.construct_scalar(node)
+    os.system("echo hello " + value)
+    return value
+
+yaml.Loader.add_constructor(u'!hello', myconstructor)
+
+trusted_yaml = "msg: !hello world"
+
+print("============ Loading trusted YAML")
+
+data = yaml.load(trusted_yaml, Loader=yaml.Loader)
+print(data)
+
+
+# somewhere else in your application
+untrusted_yaml = "msg: !hello evil"
+
+print("============ Loading untrusted YAML")
+try:
+    safe_data = yaml.safe_load(untrusted_yaml)
+    print(safe_data)
+except yaml.constructor.ConstructorError as err:
+    print("Could not load unsafe_yaml:")
+    print(format(err))
+
+

lib/yaml/__init__.py

@@ -74,7 +74,19 @@ def load(stream, Loader=Loader):
         return loader.get_single_data()
     finally:
         loader.dispose()
-safe_load = load
+
+def safe_load(stream):
+    """
+    Parse the first YAML document in a stream
+    and produce the corresponding Python object.
+
+    By default resolve only basic YAML tags
+    """
+    loader = SafeLoader(stream)
+    try:
+        return loader.get_single_data()
+    finally:
+        loader.dispose()
 
 def load_all(stream, Loader=Loader):
     """
@@ -90,23 +102,20 @@ def load_all(stream, Loader=Loader):
             yield loader.get_data()
     finally:
         loader.dispose()
-safe_load_all = load_all
-
-def danger_load(stream):
-    """
-    Parse the first YAML document in a stream
-    and produce the corresponding Python object.
-    When used on untrusted input, can result in arbitrary code execution.
-    """
-    return load(stream, DangerLoader)
 
-def danger_load_all(stream):
+def safe_load_all(stream):
     """
     Parse all YAML documents in a stream
     and produce corresponding Python objects.
-    When used on untrusted input, can result in arbitrary code execution.
+
+    By default resolve only basic YAML tags
     """
-    return load_all(stream, DangerLoader)
+    loader = SafeLoader(stream)
+    try:
+        while loader.check_data():
+            yield loader.get_data()
+    finally:
+        loader.dispose()
 
 def emit(events, stream=None, Dumper=Dumper,
         canonical=None, indent=None, width=None,
@@ -203,14 +212,6 @@ def dump_all(documents, stream=None, Dumper=Dumper,
         return getvalue()
 safe_dump_all = dump_all
 
-def danger_dump_all(documents, stream=None, **kwds):
-    """
-    Serialize a sequence of Python objects into a YAML stream.
-    Produce only basic YAML tags.
-    If stream is None, return the produced string instead.
-    """
-    return dump_all(documents, stream, Dumper=DangerDumper, **kwds)
-
 def dump(data, stream=None, Dumper=Dumper, **kwds):
     """
     Serialize a Python object into a YAML stream.
@@ -219,14 +220,6 @@ def dump(data, stream=None, Dumper=Dumper, **kwds):
     return dump_all([data], stream, Dumper=Dumper, **kwds)
 safe_dump = dump
 
-def danger_dump(data, stream=None, **kwds):
-    """
-    Serialize a Python object into a YAML stream.
-    Produce only basic YAML tags.
-    If stream is None, return the produced string instead.
-    """
-    return dump_all([data], stream, Dumper=DangerDumper, **kwds)
-
 def add_implicit_resolver(tag, regexp, first=None,
         Loader=Loader, Dumper=Dumper):
     """

lib/yaml/cyaml.py

@@ -1,6 +1,6 @@
 
-__all__ = ['CBaseLoader', 'CSafeLoader', 'CLoader', 'CDangerLoader',
-        'CBaseDumper', 'CSafeDumper', 'CDumper', 'CDangerDumper']
+__all__ = ['CBaseLoader', 'CSafeLoader', 'CLoader', 'CPythonLoader',
+        'CBaseDumper', 'CSafeDumper', 'CDumper', 'CPythonDumper']
 
 from _yaml import CParser, CEmitter
 
@@ -24,9 +24,15 @@ def __init__(self, stream):
         CParser.__init__(self, stream)
         SafeConstructor.__init__(self)
         Resolver.__init__(self)
-CSafeLoader = CLoader
 
-class CDangerLoader(CParser, Constructor, Resolver):
+class CSafeLoader(CParser, SafeConstructor, Resolver):
+
+    def __init__(self, stream):
+        CParser.__init__(self, stream)
+        SafeConstructor.__init__(self)
+        Resolver.__init__(self)
+
+class CPythonLoader(CParser, Constructor, Resolver):
 
     def __init__(self, stream):
         CParser.__init__(self, stream)
@@ -68,7 +74,7 @@ def __init__(self, stream,
         Resolver.__init__(self)
 CSafeDumper = CDumper
 
-class CDangerDumper(CEmitter, Serializer, Representer, Resolver):
+class CPythonDumper(CEmitter, Serializer, Representer, Resolver):
 
     def __init__(self, stream,
             default_style=None, default_flow_style=None,

lib/yaml/dumper.py

@@ -1,5 +1,5 @@
 
-__all__ = ['BaseDumper', 'SafeDumper', 'Dumper', 'DangerDumper']
+__all__ = ['BaseDumper', 'SafeDumper', 'Dumper', 'PythonDumper']
 
 from emitter import *
 from serializer import *
@@ -43,7 +43,7 @@ def __init__(self, stream,
         Resolver.__init__(self)
 SafeDumper = Dumper
 
-class DangerDumper(Emitter, Serializer, Representer, Resolver):
+class PythonDumper(Emitter, Serializer, Representer, Resolver):
 
     def __init__(self, stream,
             default_style=None, default_flow_style=None,

lib/yaml/loader.py

@@ -1,5 +1,5 @@
 
-__all__ = ['BaseLoader', 'SafeLoader', 'Loader', 'DangerLoader']
+__all__ = ['BaseLoader', 'SafeLoader', 'Loader', 'PythonLoader']
 
 from reader import *
 from scanner import *
@@ -27,9 +27,18 @@ def __init__(self, stream):
         Composer.__init__(self)
         SafeConstructor.__init__(self)
         Resolver.__init__(self)
-SafeLoader = Loader
 
-class DangerLoader(Reader, Scanner, Parser, Composer, Constructor, Resolver):
+class SafeLoader(Reader, Scanner, Parser, Composer, SafeConstructor, Resolver):
+
+    def __init__(self, stream):
+        Reader.__init__(self, stream)
+        Scanner.__init__(self)
+        Parser.__init__(self)
+        Composer.__init__(self)
+        SafeConstructor.__init__(self)
+        Resolver.__init__(self)
+
+class PythonLoader(Reader, Scanner, Parser, Composer, Constructor, Resolver):
 
     def __init__(self, stream):
         Reader.__init__(self, stream)

tests/lib/test_constructor.py

@@ -19,9 +19,9 @@ def _make_objects():
             NewArgs, NewArgsWithState, Reduce, ReduceWithState, MyInt, MyList, MyDict,  \
             FixedOffset, today, execute
 
-    class MyLoader(yaml.DangerLoader):
+    class MyLoader(yaml.PythonLoader):
         pass
-    class MyDumper(yaml.DangerDumper):
+    class MyDumper(yaml.PythonDumper):
         pass
 
     class MyTestClass1:

tests/lib/test_recursive.py

@@ -29,9 +29,9 @@ def test_recursive(recursive_filename, verbose=False):
     value2 = None
     output2 = None
     try:
-        output1 = yaml.danger_dump(value1)
-        value2 = yaml.danger_load(output1)
-        output2 = yaml.danger_dump(value2)
+        output1 = yaml.dump(value1, Dumper=yaml.PythonDumper)
+        value2 = yaml.load(output1, Loader=yaml.PythonLoader)
+        output2 = yaml.dump(value2, Dumper=yaml.PythonDumper)
         assert output1 == output2, (output1, output2)
     finally:
         if verbose: