From c00e26e685b7cda97fdd881c712ff4654b0cc419 Mon Sep 17 00:00:00 2001 From: d3m0n-r00t Date: Tue, 5 Jan 2021 19:41:59 +0530 Subject: [PATCH] Fixed untrusted unpickling --- petastorm/etl/legacy.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/petastorm/etl/legacy.py b/petastorm/etl/legacy.py index b16e6798c..a4422128f 100644 --- a/petastorm/etl/legacy.py +++ b/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)