FEAT Implement .skops file audit functionality

* FEAT Audit before loading a skops file

* WIP

* numpy loaders

* scipy and src issue

* sklearn

* make tests pass

* remove pickle.py

* fix a few issues

* add missing files

* add get_untrusted_types and docs

* minor fix

* add more tests

* add a smoke test, failing though

* implement safety for functions

* add missing Tree children for audit

* add missing SparseMatrixNode children for audit

* tests pass

* remove safety tree code

* minor test

* fix ids in test

* move type ignore

* add more tests and some docs

* more comments

* fix recursive dump and get_untrusted_set

* Ben's comments

* address comments: sentinel, contextmanager, sorted

* move all children to the children attribute

* add complex pipeline test

* apply Ben's suggestions

* Card object should pass trusted to load

* Update skops/io/_dispatch.py

Co-authored-by: Erin Aho <[email protected]>

Author: adrinjalali

docs/persistence.rst

@@ -49,20 +49,37 @@ The code snippet below illustrates how to use :func:`skops.io.dump` and
     clf.fit(X_train, y_train)
     dump(clf, "my-logistic-regression.skops")
     # ...
-    loaded = load("my-logistic-regression.skops")
+    loaded = load("my-logistic-regression.skops", trusted=True)
     loaded.predict(X_test)
 
     # in memory
     from skops.io import dumps, loads
     serialized = dumps(clf)
-    loaded = loads(serialized)
+    loaded = loads(serialized, trusted=True)
 
-At the moment, we support the vast majority of sklearn estimators. This includes
-complex use cases such as :class:`sklearn.pipeline.Pipeline`,
+Note that you should only load files with ``trusted=True`` if you trust the
+source. Otherwise you can get a list of untrusted types present in the dump
+using :func:`skops.io.get_untrusted_types`:
+
+.. code:: python
+
+    from skops.io import get_untrusted_types
+    unknown_types = get_untrusted_types(file="my-logistic-regression.skops")
+    print(unknown_types)
+
+Once you check the list and you validate that everything in the list is safe,
+you can load the file with ``trusted=unknown_types``:
+
+.. code:: python
+
+    loaded = load("my-logistic-regression.skops", trusted=unknown_types)
+
+At the moment, we support the vast majority of sklearn estimators. This
+includes complex use cases such as :class:`sklearn.pipeline.Pipeline`,
 :class:`sklearn.model_selection.GridSearchCV`, classes using Cython code, such
-as :class:`sklearn.tree.DecisionTreeClassifier`, and more. If you discover an sklearn
-estimator that does not work, please open an issue on the skops `GitHub page
-<https://github.com/skops-dev/skops/issues>`_ and let us know.
+as :class:`sklearn.tree.DecisionTreeClassifier`, and more. If you discover an
+sklearn estimator that does not work, please open an issue on the skops `GitHub
+page <https://github.com/skops-dev/skops/issues>`_ and let us know.
 
 In contrast to ``pickle``, skops cannot persist arbitrary Python code. This
 means if you have custom functions (say, a custom function to be used with
@@ -74,16 +91,16 @@ Roadmap
 -------
 
 Currently, it is still possible to run insecure code when using skops
-persistence. For example, it's possible to load a save file that evaluates arbitrary
-code using :func:`eval`. However, we have concrete plans on how to mitigate
-this, so please stay updated.
+persistence. For example, it's possible to load a save file that evaluates
+arbitrary code using :func:`eval`. However, we have concrete plans on how to
+mitigate this, so please stay updated.
 
 On top of trying to support persisting all relevant sklearn objects, we plan on
-making persistence extensible for other libraries. As a user, this means that if
-you trust a certain library, you will be able to tell skops to load code from
-that library. As a library author, there will be a clear path of what needs to
-be done to add secure persistence to your library, such that skops can save and
-load code from your library.
+making persistence extensible for other libraries. As a user, this means that
+if you trust a certain library, you will be able to tell skops to load code
+from that library. As a library author, there will be a clear path of what
+needs to be done to add secure persistence to your library, such that skops can
+save and load code from your library.
 
 To follow what features are currently planned, filter for the `"persistence"
 label <https://github.com/skops-dev/skops/labels/persistence>`_ in our GitHub

skops/card/_model_card.py

@@ -165,20 +165,25 @@ def metadata_from_config(config_path: Union[str, Path]) -> ModelCardData:
     return card_data
 
 
-def _load_model(model: Any) -> Any:
-    """Loads the mddel if provided a file path, if already a model instance return it
-    unmodified.
+def _load_model(model: Any, trusted=False) -> Any:
+    """Return a model instance.
+
+    Loads the model if provided a file path, if already a model instance return
+    it unmodified.
 
     Parameters
     ----------
     model : pathlib.Path, str, or sklearn estimator
         Path/str or the actual model instance. if a Path or str, loads the model.
 
+    trusted : bool, default=False
+        Passed to :func:`skops.io.load` if the model is a file path and it's
+        a `skops` file.
+
     Returns
     -------
     model : object
         Model instance.
-
     """
 
     if not isinstance(model, (Path, str)):
@@ -190,11 +195,11 @@ def _load_model(model: Any) -> Any:
 
     try:
         if zipfile.is_zipfile(model_path):
-            model = load(model_path)
+            model = load(model_path, trusted=trusted)
         else:
             model = joblib.load(model_path)
     except Exception as ex:
-        msg = f'An "{type(ex).__name__}" occured during model loading.'
+        msg = f'An "{type(ex).__name__}" occurred during model loading.'
         raise RuntimeError(msg) from ex
 
     return model
@@ -227,6 +232,10 @@ class Card:
         of the ``config.json`` file, which itself is created by
         :func:`skops.hub_utils.init`.
 
+    trusted: bool, default=False
+        Passed to :func:`skops.io.load` if the model is a file path and it's
+        a `skops` file.
+
     Attributes
     ----------
     model: estimator object
@@ -294,13 +303,15 @@ def __init__(
         model: Any,
         model_diagram: bool = True,
         metadata: Optional[ModelCardData] = None,
+        trusted: bool = False,
     ) -> None:
         self.model = model
         self.model_diagram = model_diagram
         self._eval_results = {}  # type: ignore
         self._template_sections: dict[str, str] = {}
         self._extra_sections: list[tuple[str, Any]] = []
         self.metadata = metadata or ModelCardData()
+        self.trusted = trusted
 
     def get_model(self) -> Any:
         """Returns sklearn estimator object if ``Path``/``str``
@@ -311,7 +322,7 @@ def get_model(self) -> Any:
         model : Object
             Model instance.
         """
-        model = _load_model(self.model)
+        model = _load_model(self.model, self.trusted)
         return model
 
     def add(self, **kwargs: str) -> "Card":

skops/card/tests/test_card.py

@@ -40,10 +40,10 @@ def save_model_to_file(model_instance, suffix):
 def test_load_model(suffix):
     model0 = LinearRegression(n_jobs=123)
     _, save_file = save_model_to_file(model0, suffix)
-    loaded_model_str = _load_model(save_file)
+    loaded_model_str = _load_model(save_file, trusted=True)
     save_file_path = Path(save_file)
-    loaded_model_path = _load_model(save_file_path)
-    loaded_model_instance = _load_model(model0)
+    loaded_model_path = _load_model(save_file_path, trusted=True)
+    loaded_model_instance = _load_model(model0, trusted=True)
 
     assert loaded_model_str.n_jobs == 123
     assert loaded_model_path.n_jobs == 123
@@ -431,7 +431,7 @@ def test_with_metadata(self, card: Card, meth):
 
 class TestCardModelAttribute:
     def path_to_card(self, path):
-        card = Card(model=path)
+        card = Card(model=path, trusted=True)
         card.add(
             model_description="A description",
             model_card_authors="Jane Doe",
@@ -470,7 +470,7 @@ def test_load_model_exception(self, meth, suffix):
 
         os.close(file_handle)
 
-        with pytest.raises(Exception, match="occured during model loading."):
+        with pytest.raises(Exception, match="occurred during model loading."):
             card = Card(file_name)
             meth(card)
 

skops/io/__init__.py

@@ -1,3 +1,3 @@
-from ._persist import dump, dumps, load, loads
+from ._persist import dump, dumps, get_untrusted_types, load, loads
 
-__all__ = ["dumps", "load", "loads", "dump"]
+__all__ = ["dumps", "load", "loads", "dump", "get_untrusted_types"]

skops/io/_audit.py

@@ -0,0 +1,61 @@
+from skops.io.exceptions import UntrustedTypesFoundException
+
+
+def check_type(module_name, type_name, trusted):
+    """Check if a type is safe to load.
+
+    A type is safe to load only if it's present in the trusted list.
+
+    Parameters
+    ----------
+    module_name : str
+        The module name of the type.
+
+    type_name : str
+        The class name of the type.
+
+    trusted : bool, or list of str
+        If ``True``, the tree is considered safe. Otherwise trusted has to be
+        a list of trusted types.
+
+    Returns
+    -------
+    is_safe : bool
+        True if the type is safe, False otherwise.
+    """
+    if trusted is True:
+        return True
+    return module_name + "." + type_name in trusted
+
+
+def audit_tree(tree, trusted):
+    """Audit a tree of nodes.
+
+    A tree is safe if it only contains trusted types. Audit is skipped if
+    trusted is ``True``.
+
+    Parameters
+    ----------
+    tree : skops.io._dispatch.Node
+        The tree to audit.
+
+    trusted : bool, or list of str
+        If ``True``, the tree is considered safe. Otherwise trusted has to be
+        a list of trusted types names.
+
+        An entry in the list is typically of the form
+        ``skops.io._utils.get_module(obj) + "." + obj.__class__.__name__``.
+
+    Raises
+    ------
+    UntrustedTypesFoundException
+        If the tree contains an untrusted type.
+    """
+    if trusted is True:
+        return
+
+    unsafe = tree.get_unsafe_set()
+    if isinstance(trusted, (list, set)):
+        unsafe -= set(trusted)
+    if unsafe:
+        raise UntrustedTypesFoundException(unsafe)