SEC fix an issue with skops update where it'd load any types

docs/changes.rst

@@ -11,6 +11,12 @@ skops Changelog
 
 v0.15
 -----
+- Fix a security issue where ``skops update`` (:mod:`skops.cli`) would load a
+  file while blindly trusting every type declared in it, allowing a malicious
+  ``.skops`` file to execute arbitrary code. ``skops update`` now refuses to
+  load types that are not trusted by default; use the new ``--trusted`` option
+  to explicitly allow types you have reviewed.
+  :pr:`533` by `Adrin Jalali`_.
 - Support persisting scipy sparse *arrays* (e.g. ``csr_array``) through the same
   efficient ``npz`` format used for sparse matrices, and handle the
   ``scipy.special`` ufunc wrappers introduced in scipy 2.0. This adds

docs/persistence.rst

@@ -164,6 +164,19 @@ The below command is an example on how to create an updated version of a file
 
     skops update my_model.skops -o my_model-updated.skops
 
+Updating a file requires loading it, and loading is only safe for trusted types.
+By default ``skops update`` only loads types that skops trusts, and refuses to
+proceed if the file contains other types, printing the offending types. If, after
+reviewing them, you trust those types, pass them explicitly with ``--trusted``:
+
+.. code-block:: console
+
+    skops update my_model.skops -o my_model-updated.skops --trusted a.Type b.OtherType
+
+Never pass types to ``--trusted`` without reviewing them first: a malicious file
+can declare arbitrary types, and trusting them can lead to arbitrary code
+execution on load.
+
 Further help for the different supported options can be found by calling
 ``skops update --help`` in a terminal.
 

skops/cli/_update.py

@@ -11,12 +11,14 @@
 from skops.cli._utils import get_log_level
 from skops.io import dump, get_untrusted_types, load
 from skops.io._protocol import PROTOCOL
+from skops.io.exceptions import UntrustedTypesFoundException
 
 
 def _update_file(
     input_file: str | Path,
     output_file: str | Path | None = None,
     inplace: bool = False,
+    trusted: list[str] | None = None,
     logger: logging.Logger = logging.getLogger(),
 ) -> None:
     """Function that is called by ``skops update`` entrypoint.
@@ -35,6 +37,14 @@ def _update_file(
     inplace : bool, default=False
         Whether to update and overwrite the input file in place.
 
+    trusted : list of str, default=None
+        List of types the caller has reviewed and trusts to be loaded from the
+        input file, in addition to the types skops trusts by default. Updating a
+        file requires loading it, and loading is only safe for trusted types.
+        ``skops update`` refuses to load a file containing types outside this
+        list, since blindly trusting every type declared in the file would let a
+        malicious file execute arbitrary code on load.
+
     logger : logging.Logger, default=logging.getLogger()
         Logger to use for logging.
     """
@@ -48,7 +58,18 @@ def _update_file(
                 " file."
             )
 
-    input_model = load(input_file, trusted=get_untrusted_types(file=input_file))
+    trusted = list(trusted) if trusted is not None else []
+    unreviewed = [t for t in get_untrusted_types(file=input_file) if t not in trusted]
+    if unreviewed:
+        logger.error(
+            "The input file contains types that are not trusted by default: "
+            f"{unreviewed}. Updating a file requires loading it, and skops will "
+            "not load untrusted types automatically. Review these types and, if "
+            "you trust them, re-run with `--trusted` listing each one explicitly."
+        )
+        raise UntrustedTypesFoundException(unreviewed)
+
+    input_model = load(input_file, trusted=trusted)
     with zipfile.ZipFile(input_file, "r") as zip_file:
         input_file_schema = json.loads(zip_file.read("schema.json"))
 
@@ -103,6 +124,18 @@ def format_parser(
         help="Update and overwrite the input file in place.",
         action="store_true",
     )
+    parser_subgroup.add_argument(
+        "--trusted",
+        nargs="*",
+        default=[],
+        help=(
+            "Types to trust when loading the input file, in addition to the types "
+            "trusted by default. Updating requires loading the file, and skops "
+            "refuses to load types it does not trust. Inspect the file's untrusted "
+            "types first (e.g. with `skops.io.get_untrusted_types`), and only list "
+            "types here after you have reviewed them."
+        ),
+    )
     parser_subgroup.add_argument(
         "-v",
         "--verbose",
@@ -124,6 +157,7 @@ def main(
     output_file = Path(parsed_args.output_file) if parsed_args.output_file else None
     input_file = Path(parsed_args.input)
     inplace = parsed_args.inplace
+    trusted = parsed_args.trusted
 
     logging.basicConfig(format="%(levelname)-8s: %(message)s")
     logger.setLevel(level=get_log_level(parsed_args.loglevel))
@@ -132,5 +166,6 @@ def main(
         input_file=input_file,
         output_file=output_file,
         inplace=inplace,
+        trusted=trusted,
         logger=logger,
     )

skops/cli/tests/test_entrypoint.py

@@ -57,5 +57,6 @@ def test_update_works_as_expected(
             input_file=pathlib.Path("abc.skops"),
             output_file=pathlib.Path("abc-new.skops"),
             inplace=False,
+            trusted=[],
             logger=mock.ANY,
         )

skops/cli/tests/test_update.py

@@ -1,5 +1,8 @@
+import io
+import json
 import logging
 import pathlib
+import zipfile
 from functools import partial
 from unittest import mock
 
@@ -8,6 +11,7 @@
 
 from skops.cli import _update
 from skops.io import _persist, _protocol, dump, load
+from skops.io.exceptions import UntrustedTypesFoundException
 
 
 class TestUpdate:
@@ -150,6 +154,76 @@ def test_error_with_output_file_and_inplace(
                 input_file=skops_path, output_file=new_skops_path, inplace=True
             )
 
+    @pytest.fixture
+    def malicious_path(self, tmp_path: pathlib.Path) -> pathlib.Path:
+        """A crafted skops file that tries to run ``os.system`` via a LossNode.
+
+        See the CVE-style report: an attacker can declare an arbitrary callable
+        as a LossNode constructor. ``get_untrusted_types`` surfaces it (here as
+        ``os.system``), and blindly trusting that output would execute it.
+        """
+        schema = {
+            "__class__": "system",
+            "__module__": "os",
+            "__loader__": "LossNode",
+            "__reduce__": {
+                "args": {
+                    "__class__": "tuple",
+                    "__module__": "builtins",
+                    "__loader__": "TupleNode",
+                    "content": [
+                        {
+                            "__class__": "str",
+                            "__module__": "builtins",
+                            "__loader__": "JsonNode",
+                            "content": '"echo pwned"',
+                            "is_json": True,
+                        }
+                    ],
+                }
+            },
+            "content": {
+                "__class__": "dict",
+                "__module__": "builtins",
+                "__loader__": "DictNode",
+                "content": {},
+                "key_types": {
+                    "__class__": "list",
+                    "__module__": "builtins",
+                    "__loader__": "ListNode",
+                    "content": [],
+                },
+            },
+            "protocol": _protocol.PROTOCOL,
+            "_skops_version": "0.14.0",
+        }
+        path = tmp_path / "malicious.skops"
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as zf:
+            zf.writestr("schema.json", json.dumps(schema))
+        path.write_bytes(buf.getvalue())
+        return path
+
+    def test_refuses_untrusted_types_by_default(
+        self,
+        malicious_path: pathlib.Path,
+        new_skops_path: pathlib.Path,
+    ):
+        """``skops update`` must not auto-trust the types declared in the file.
+
+        This is the core of the RCE fix: previously the CLI fed
+        ``get_untrusted_types()`` straight back into ``load()``, so a file could
+        smuggle any type (e.g. ``os.system``) into the trusted set with no human
+        review. Now untrusted types are refused unless the operator lists them
+        explicitly via ``--trusted``.
+        """
+        with pytest.raises(UntrustedTypesFoundException, match="os.system"):
+            _update._update_file(
+                input_file=malicious_path,
+                output_file=new_skops_path,
+            )
+        assert not new_skops_path.exists()
+
 
 class TestMain:
     @pytest.fixture
@@ -172,6 +246,7 @@ def test_output_argument(
             input_file=pathlib.Path(input_path),
             output_file=pathlib.Path(output_path),
             inplace=False,
+            trusted=[],
             logger=tmp_logger,
         )
 
@@ -189,6 +264,26 @@ def test_inplace_argument(
             input_file=pathlib.Path(input_path),
             output_file=None,
             inplace=True,
+            trusted=[],
+            logger=tmp_logger,
+        )
+
+    @mock.patch("skops.cli._update._update_file")
+    def test_trusted_argument(
+        self, mock_update: mock.MagicMock, tmp_logger: logging.Logger
+    ):
+        input_path = "abc.skops"
+        output_path = "abc-new.skops"
+        namespace, _ = _update.format_parser().parse_known_args(
+            [input_path, "-o", output_path, "--trusted", "foo.Bar", "baz.Qux"]
+        )
+
+        _update.main(namespace, tmp_logger)
+        mock_update.assert_called_once_with(
+            input_file=pathlib.Path(input_path),
+            output_file=pathlib.Path(output_path),
+            inplace=False,
+            trusted=["foo.Bar", "baz.Qux"],
             logger=tmp_logger,
         )
 
@@ -222,6 +317,7 @@ def test_given_log_levels_works_as_expected(
             input_file=pathlib.Path(input_path),
             output_file=pathlib.Path(output_path),
             inplace=False,
+            trusted=[],
             logger=tmp_logger,
         )
         assert tmp_logger.getEffectiveLevel() == expected_level