diff --git a/requirements.txt b/requirements.txt index 06919ef..60bd7e4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ conference-scheduler>=2.1.3 daiquiri>=1.2.1 python-slugify>=1.2.4 ruamel.yaml>=0.15.23 +future>=0.18.2 \ No newline at end of file diff --git a/src/scheduler/io.py b/src/scheduler/io.py index e2cb8c7..f10459d 100644 --- a/src/scheduler/io.py +++ b/src/scheduler/io.py @@ -1,4 +1,6 @@ +import io import csv +import builtins import pickle from pathlib import Path from pprint import pformat @@ -13,6 +15,28 @@ yaml.default_flow_style = False +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() + def import_yaml(input_folder): """Import all yaml files in the given folder into a single resources dict""" @@ -69,6 +93,8 @@ def import_schedule_definition(solution_folder): logger.info( f'Importing resources, events, slots and schedule from {pickle_file}') with pickle_file.open('rb') as f: + """restricting unsafe pickle files""" + restricted_loads(f.read()) bundle = pickle.load(f) return bundle