Merge pull request #4 from yilun11/CVE-2020-13091

CVE 2020 13091

Author: rvema

INFO

@@ -0,0 +1,30 @@
+### tests
+
+updated /tests/io/test_pickle.py
+did not write the compatibility tests for older pandas versions yet
+
+###
+
+# call this file pickle_config.yml
+# this is meant to configure pickle deserialization, which can be a security risk
+ 
+# mode can be permit, deny, off
+ 
+# permit means pickle is only allowed to load permitted libraries (safer option)
+# deny means pickle is not allowed to load denied libraries (less safe option)
+# off means no security and should only be used on safe pickles (least safe option)
+mode: permit
+ 
+# here you would list packages and classes that can be used
+# this is how Python recommend you restrict pickling globals
+# see https://docs.python.org/3/library/pickle.html#restricting-globals
+permit:
+  builtins: ['range', 'complex', 'set', 'frozenset', 'slice']
+ 
+# here you would list packages and classes than cannot be used
+# all others will be allowed if running on deny mode
+# this is not as safe as the permit method
+# see https://pythonmana.com/2022/143/202205231222219535.html
+deny:
+  builtins: ['eval', 'exec', 'execfile', 'compile', 'open', 'input', '__import__', 'exit']
+  os: ['system']

pandas/compat/pickle_compat.py

@@ -15,6 +15,8 @@
 
 import numpy as np
 
+from pandas._config import get_option
+
 from pandas._libs.arrays import NDArrayBacked
 from pandas._libs.tslibs import BaseOffset
 
@@ -196,14 +198,28 @@ def __new__(cls) -> DataFrame:  # type: ignore[misc]
 
 # our Unpickler sub-class to override methods and some dispatcher
 # functions for compat and uses a non-public class of the pickle module.
+# checks modules against permit/deny list and raises error if module is not forbidden.
 
 
 class Unpickler(pkl._Unpickler):
     def find_class(self, module, name):
         # override superclass
         key = (module, name)
         module, name = _class_locations_map.get(key, key)
-        return super().find_class(module, name)
+        opt = get_option("pickler.unpickle.mode")
+        # Only allow safe modules and classes. Tuples defined in config
+        # Do not allow unsafe modules and classes.
+        if (
+            (opt == "off")
+            or (opt == "permit" and (module, name) in get_option("pickler.safe.tuples"))
+            or (
+                opt == "deny"
+                and (module, name) not in get_option("pickler.unsafe.tuples")
+            )
+        ):
+            return super().find_class(module, name)
+        # Forbid everything else.
+        raise pkl.UnpicklingError(f"global '{module} . {name}' is forbidden")
 
 
 Unpickler.dispatch = copy.copy(Unpickler.dispatch)

pandas/core/config_init.py

@@ -16,6 +16,8 @@
 from typing import Callable
 import warnings
 
+import yaml
+
 import pandas._config.config as cf
 from pandas._config.config import (
     is_bool,
@@ -970,3 +972,78 @@ def register_converter_cb(key) -> None:
         styler_environment,
         validator=is_instance_factory([type(None), str]),
     )
+
+# ------
+# Pickler
+# ------
+
+pickler_unpickle_mode = """
+: str
+    Determine which mode to use in {"off", "permit", "deny"}.
+"""
+
+pickler_safe_tuples = """
+: array
+    Used when pickler.unpickle.mode is "permit"
+    Array of safe tuples, e.g. [("builtins", "range"), ("builtins", "complex")]
+"""
+
+pickler_unsafe_tuples = """
+: array
+    Used when pickler.unpickle.mode is "deny"
+    Array of unsafe tuples, e.g. [("os", "system"), ("joblib", "load")]
+"""
+
+# location of config file from env var
+str_loc = os.environ.get("PANDAS_UNPICKLE_SECURE", "pickle_config.yml")
+safe_tuples = []
+unsafe_tuples = []
+
+if os.path.exists(str_loc):
+    pickle_config = yaml.load(open(str_loc), Loader=yaml.SafeLoader)
+
+    if pickle_config["mode"] == "permit":
+        for k, v in pickle_config["permit"].items():
+            for i in v:
+                safe_tuples.append((k, i))
+    elif pickle_config["mode"] == "deny":
+        for k, v in pickle_config["deny"].items():
+            for i in v:
+                unsafe_tuples.append((k, i))
+else:
+    pickle_config = {}
+    pickle_config["mode"] = "deny"
+    # see deny list example at https://pythonmana.com/2022/143/202205231222219535.html
+    unsafe_tuples = [
+        ("os", "system"),
+        ("posix", "system"),
+        ("builtins", "eval"),
+        ("builtins", "exec"),
+        ("builtins", "execfile"),
+        ("builtins", "compile"),
+        ("builtins", "open"),
+        ("builtins", "import"),
+        ("builtins", "__import__"),
+        ("builtins", "exit"),
+    ]
+
+with cf.config_prefix("pickler"):
+    cf.register_option(
+        "unpickle.mode",
+        # get the default value from the config file
+        pickle_config["mode"],
+        pickler_unpickle_mode,
+        validator=is_one_of_factory(["off", "permit", "deny"]),
+    )
+
+    cf.register_option(
+        "safe.tuples",
+        safe_tuples,
+        pickler_safe_tuples,
+    )
+
+    cf.register_option(
+        "unsafe.tuples",
+        unsafe_tuples,
+        pickler_unsafe_tuples,
+    )

pandas/io/pickle.py

@@ -5,6 +5,8 @@
 from typing import Any
 import warnings
 
+from pandas._config import get_option
+
 from pandas._typing import (
     CompressionOptions,
     FilePath,
@@ -101,15 +103,8 @@ def to_pickle(
         is_text=False,
         storage_options=storage_options,
     ) as handles:
-        if handles.compression["method"] in ("bz2", "xz") and protocol >= 5:
-            # some weird TypeError GH#39002 with pickle 5: fallback to letting
-            # pickle create the entire object and then write it to the buffer.
-            # "zip" would also be here if pandas.io.common._BytesZipFile
-            # wouldn't buffer write calls
-            handles.handle.write(pickle.dumps(obj, protocol=protocol))
-        else:
-            # letting pickle write directly to the buffer is more memory-efficient
-            pickle.dump(obj, handles.handle, protocol=protocol)
+        # letting pickle write directly to the buffer is more memory-efficient
+        pickle.dump(obj, handles.handle, protocol=protocol)
 
 
 @doc(
@@ -122,12 +117,15 @@ def read_pickle(
     storage_options: StorageOptions = None,
 ):
     """
-    Load pickled pandas object (or any object) from file.
+    Load pickled pandas object (or any object) from file. By default, only a
+    safe subset of classes from builtins can be called while loading the
+    object. See INFO file for customizing the security settings.
 
     .. warning::
 
        Loading pickled data received from untrusted sources can be
-       unsafe. See `here <https://docs.python.org/3/library/pickle.html>`__.
+       unsafe if not using the default security settings.
+       See `here <https://docs.python.org/3/library/pickle.html>`__.
 
     Parameters
     ----------
@@ -186,6 +184,27 @@ def read_pickle(
     3    3    8
     4    4    9
     """
+    
+    class RestrictedUnpickler(pickle.Unpickler):
+        def find_class(self, module, name):
+            opt = get_option("pickler.unpickle.mode")
+            # Only allow safe modules and classes. Tuples defined in config
+            # Do not allow unsafe modules and classes.
+            if (
+                (opt == "off")
+                or (
+                    opt == "permit"
+                    and (module, name) in get_option("pickler.safe.tuples")
+                )
+                or (
+                    opt == "deny"
+                    and (module, name) not in get_option("pickler.unsafe.tuples")
+                )
+            ):
+                return super().find_class(module, name)
+            # Forbid everything else.
+            raise pickle.UnpicklingError(f"global '{module} . {name}' is forbidden")
+
     excs_to_catch = (AttributeError, ImportError, ModuleNotFoundError, TypeError)
     with get_handle(
         filepath_or_buffer,
@@ -205,7 +224,7 @@ def read_pickle(
                 with warnings.catch_warnings(record=True):
                     # We want to silence any warnings about, e.g. moved modules.
                     warnings.simplefilter("ignore", Warning)
-                    return pickle.load(handles.handle)
+                    return RestrictedUnpickler.load(handles.handle)
             except excs_to_catch:
                 # e.g.
                 #  "No module named 'pandas.core.sparse.series'"

pandas/tests/io/data/pickle/test_forbidden.pkl

@@ -397,5 +397,10 @@ def test_round_trip_valid_encodings(self, enc, df):
     )
     def test_raw_roundtrip(self, data):
         # PR #25040 wide unicode wasn't copied correctly on PY3 on windows
-        clipboard_set(data)
-        assert data == clipboard_get()
+        # adding coverage for when not implemented
+        try:
+            clipboard_set(data)
+            assert data == clipboard_get()
+        except PyperclipException:
+            with pytest.raises(PyperclipException, match=r".*not-implemented-error.*"):
+                clipboard_set(data)

pandas/tests/io/test_clipboard.py

@@ -598,3 +598,20 @@ def test_pickle_frame_v124_unpickle_130():
 
     expected = pd.DataFrame()
     tm.assert_frame_equal(df, expected)
+
+
+def test_read_pickle_forbidden():
+    # related to CVE - https://nvd.nist.gov/vuln/detail/CVE-2020-13091
+
+    class MyEvilPickle:
+        def __reduce__(self):
+            return (os.system, ("whoami",))
+
+    pickle_data = pickle.dumps(MyEvilPickle())
+    # storing the serialized output into a file in current directory
+    path = os.path.join(os.path.dirname(__file__), "data", "pickle", "test_forbidden.pkl")
+    with open(path, "wb") as file:
+        file.write(pickle_data)
+
+    with pytest.raises(pickle.UnpicklingError, match=r".* forbidden"):
+        pd.read_pickle(path)