Security Fix for Arbitrary Code Execution - huntr.dev

petastorm/etl/legacy.py

@@ -16,6 +16,32 @@
 
 from six.moves import cPickle as pickle
 
+import io
+import builtins
+
+safe_builtins = {
+    'range',
+    'complex',
+    'set',
+    'frozenset',
+    'slice',
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+
+    def find_class(self, module, name):
+        """Only allow safe classes from builtins"""
+        if module == "builtins" and name in safe_builtins:
+            return getattr(builtins, name)
+        """Forbid everything else"""
+        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
+                                     (module, name))
+
+def restricted_loads(s):
+    """Helper function analogous to pickle.loads()"""
+    return RestrictedUnpickler(io.BytesIO(s)).load()
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,5 +69,5 @@ def depickle_legacy_package_name_compatible(pickled_string):
                                'Regenerate metadata.', legacy_package_name, legacy_module, legacy_module)
 
             pickled_string = modified_pickled_string
-
+    restricted_loads(pickled_string)
     return pickle.loads(pickled_string)