CVE-2024-2083 (Python, published 2024-04-16T00:15Z)
Description: value: A directory traversal vulnerability exists in the zenml-io/zenml repository, specifically within the /api/v1/steps endpoint. Attackers can exploit this vulnerability by manipulating the 'logs' URI path in the request to fetch arbitrary file content, bypassing intended access restrictions. The vulnerability arises due to the lack of validation for directory traversal patterns, allowing attackers to access files outside of the restricted directory.
Severity: nan (CVSS v3 base score: )
Weakness: NVD-CWE-noinfo \u2014 Insufficient Information
There is insufficient information about the issue to classify it; details are unkown or unspecified.
Repository: https://github.com/zenml-io/zenml
Files changed:
src/zenml/artifact_stores/base_artifact_store.py
src/zenml/artifacts/utils.py
src/zenml/logging/step_logging.py
src/zenml/materializers/base_materializer.py
src/zenml/materializers/built_in_materializer.py
src/zenml/materializers/cloudpickle_materializer.py
src/zenml/materializers/numpy_materializer.py
src/zenml/materializers/pandas_materializer.py
src/zenml/materializers/service_materializer.py
src/zenml/materializers/structured_string_materializer.py
src/zenml/orchestrators/output_utils.py
tests/integration/functional/artifacts/test_base_artifact_store.py
tests/integration/functional/artifacts/test_utils.py
tests/unit/artifacts/test_utils.py
tests/unit/materializers/test_built_in_materializer.py
tests/unit/materializers/test_cloudpickle_materializer.py
tests/unit/test_general.py
Commit message: Improve Artifact Store isolation (#2490)

* dir traversal issue

* Auto-update of Starter template

* Auto-update of NLP template

* reroute artifacts and logs via AS

* reroute materializers via AS

* simplify to one deco

* fix materializer tests

* allow local download

* Auto-update of E2E template

* fix test issues

* rework based on comments

* fix bugs

* lint

* Candidate (#2493)

Co-authored-by: Stefan Nica <stefan@zenml.io>

* darglint

---------

Co-authored-by: GitHub Actions <actions@github.com>
Co-authored-by: Stefan Nica <stefan@zenml.io>

Vulnerable code:
def _sanitize_potential_path(potential_path: Any) -> Any:
    """Sanitizes the input if it is a path.
    If the input is a **remote** path, this function replaces backslash path
    separators by forward slashes.
        potential_path: Value that potentially refers to a (remote) path.
        The original input or a sanitized version of it in case of a remote
        path.
    if isinstance(potential_path, bytes):
        path = fileio.convert_to_str(potential_path)
    elif isinstance(potential_path, str):
        path = potential_path
    else:
        # Neither string nor bytes, this is not a path
        return potential_path
    if io_utils.is_remote(path):
        # If we have a remote path, replace windows path separators with
        # slashes
        import ntpath
        import posixpath
        path = path.replace(ntpath.sep, posixpath.sep)
    return path
def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
    """Sanitizes path inputs before calling the original function.
    Args:
        _func: The function for which to sanitize the inputs.
    Returns:
        Function that calls the input function with sanitized path inputs.
    """
    def inner_function(*args: Any, **kwargs: Any) -> Any:
        """Inner function.
        args = tuple(_sanitize_potential_path(arg) for arg in args)
            key: _sanitize_potential_path(value)
        return _func(*args, **kwargs)

    return inner_function
            self.__class__.__name__,
            (BaseFilesystem,),
            {
                "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
                "open": staticmethod(_sanitize_paths(self.open)),
                "copyfile": staticmethod(_sanitize_paths(self.copyfile)),
                "exists": staticmethod(_sanitize_paths(self.exists)),
                "glob": staticmethod(_sanitize_paths(self.glob)),
                "isdir": staticmethod(_sanitize_paths(self.isdir)),
                "listdir": staticmethod(_sanitize_paths(self.listdir)),
                "makedirs": staticmethod(_sanitize_paths(self.makedirs)),
                "mkdir": staticmethod(_sanitize_paths(self.mkdir)),
                "remove": staticmethod(_sanitize_paths(self.remove)),
                "rename": staticmethod(_sanitize_paths(self.rename)),
                "rmtree": staticmethod(_sanitize_paths(self.rmtree)),
                "size": staticmethod(_sanitize_paths(self.size)),
                "stat": staticmethod(_sanitize_paths(self.stat)),
                "walk": staticmethod(_sanitize_paths(self.walk)),
            },
    if manual_save and fileio.exists(uri):
    fileio.makedirs(uri)
    with fileio.open(
from zenml.io import fileio
    if not fileio.exists(logs_base_uri):
        fileio.makedirs(logs_base_uri)
    if fileio.exists(logs_uri):
        fileio.remove(logs_uri)
                    with fileio.open(self.logs_uri, "a") as file:
        with fileio.open(visualization_uri, "w") as f:
from zenml.io import fileio
        with fileio.open(self.data_path, "rb") as file_:
        with fileio.open(self.data_path, "wb") as file_:
        if not fileio.exists(self.data_path) and not fileio.exists(
            self.metadata_path
        ):
        if fileio.exists(self.data_path):
                fileio.mkdir(element_path)
            if fileio.exists(self.metadata_path):
                fileio.remove(self.metadata_path)
                fileio.rmtree(entry["path"])
from zenml.io import fileio
        with fileio.open(filepath, "rb") as fid:
        with fileio.open(filepath, "wb") as fid:
from zenml.io import fileio
        if fileio.exists(numpy_file):
            with fileio.open(numpy_file, "rb") as f:
        elif fileio.exists(os.path.join(self.uri, DATA_FILENAME)):
                with fileio.open(
        with fileio.open(os.path.join(self.uri, NUMPY_FILENAME), "wb") as f:
        with fileio.open(output_path, "wb") as f:
        with fileio.open(output_path, "wb") as f:
from zenml.io import fileio
        if fileio.exists(self.parquet_path):
                with fileio.open(self.parquet_path, mode="rb") as f:
            with fileio.open(self.csv_path, mode="rb") as f:
            with fileio.open(self.parquet_path, mode="wb") as f:
            with fileio.open(self.csv_path, mode="wb") as f:
        with fileio.open(describe_uri, mode="wb") as f:
from zenml.io import fileio
        with fileio.open(filepath, "r") as f:
        with fileio.open(filepath, "w") as f:
from zenml.io import fileio
        with fileio.open(self._get_filepath(data_type), "r") as f:
        with fileio.open(self._get_filepath(type(data)), "w") as f:
from zenml.io import fileio
        if fileio.exists(artifact_uri):
        fileio.makedirs(artifact_uri)
        if fileio.isdir(artifact_uri):
            fileio.rmtree(artifact_uri)
    tmp_path, clean_client_with_run
def model_metadata_dir(model_artifact):
    temp_dir = tempfile.mkdtemp()
def numpy_file_uri():
    temp_dir = tempfile.mkdtemp()
def test_container_materializer_for_custom_types(mocker):
    with TemporaryDirectory() as artifact_uri:
    with TemporaryDirectory() as artifact_uri:
    with TemporaryDirectory() as artifact_uri:

Fixed code:
import inspect
from pathlib import Path
class _sanitize_paths:
    """Sanitizes path inputs before calling the original function.
    Extra decoration layer is needed to pass in fixed artifact store root
    path for static methods that are called on filesystems directly.
        func: The function to decorate.
        fixed_root_path: The fixed artifact store root path.
        is_static: Whether the function is static or not.
        Function that calls the input function with sanitized path inputs.
    def __init__(self, func: Callable[..., Any], fixed_root_path: str) -> None:
        """Initializes the decorator.
        Args:
            func: The function to decorate.
            fixed_root_path: The fixed artifact store root path.
        """
        self.func = func
        self.fixed_root_path = fixed_root_path
        self.path_args: List[int] = []
        self.path_kwargs: List[str] = []
        for i, param in enumerate(
            inspect.signature(self.func).parameters.values()
        ):
            if param.annotation == PathType:
                self.path_kwargs.append(param.name)
                if param.default == inspect.Parameter.empty:
                    self.path_args.append(i)
    def _validate_path(self, path: str) -> None:
        """Validates a path.
        Args:
            path: The path to validate.
        Raises:
            FileNotFoundError: If the path is outside of the artifact store
                bounds.
        """
        if not path.startswith(self.fixed_root_path):
            raise FileNotFoundError(
                f"File `{path}` is outside of "
                f"artifact store bounds `{self.fixed_root_path}`"
            )
    def _sanitize_potential_path(self, potential_path: Any) -> Any:
        """Sanitizes the input if it is a path.

        If the input is a **remote** path, this function replaces backslash path
        separators by forward slashes.
        Args:
            potential_path: Value that potentially refers to a (remote) path.

        Returns:
            The original input or a sanitized version of it in case of a remote
            path.
        """
        if isinstance(potential_path, bytes):
            path = fileio.convert_to_str(potential_path)
        elif isinstance(potential_path, str):
            path = potential_path
        else:
            # Neither string nor bytes, this is not a path
            return potential_path

        if io_utils.is_remote(path):
            # If we have a remote path, replace windows path separators with
            # slashes
            import ntpath
            import posixpath

            path = path.replace(ntpath.sep, posixpath.sep)
            self._validate_path(path)
        else:
            self._validate_path(str(Path(path).absolute().resolve()))

        return path

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Decorator function that sanitizes paths before calling the original function.
        # verify if `self` is part of the args
        has_self = bool(args and isinstance(args[0], BaseArtifactStore))

        # sanitize inputs for relevant args and kwargs, keep rest unchanged
        args = tuple(
            self._sanitize_potential_path(
                arg,
            )
            if i + has_self in self.path_args
            else arg
            for i, arg in enumerate(args)
        )
            key: self._sanitize_potential_path(
                value,
            )
            if key in self.path_kwargs
            else value
        return self.func(*args, **kwargs)
    @abstractmethod
        overloads: Dict[str, Any] = {
            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
        }
        for abc_method in inspect.getmembers(BaseArtifactStore):
            if getattr(abc_method[1], "__isabstractmethod__", False):
                sanitized_method = _sanitize_paths(
                    getattr(self, abc_method[0]), self.path
                )
                # prepare overloads for filesystem methods
                overloads[abc_method[0]] = staticmethod(sanitized_method)

                # decorate artifact store methods
                setattr(
                    self,
                    abc_method[0],
                    sanitized_method,
                )

            self.__class__.__name__, (BaseFilesystem,), overloads
    if manual_save and artifact_store.exists(uri):
    artifact_store.makedirs(uri)
        IOError: If the artifact store rejects the request.
    except IOError as e:
        raise e
    artifact_store = Client().active_stack.artifact_store
    with artifact_store.open(
from zenml.client import Client
    artifact_store = Client().active_stack.artifact_store
    if not artifact_store.exists(logs_base_uri):
        artifact_store.makedirs(logs_base_uri)
    if artifact_store.exists(logs_uri):
        artifact_store.remove(logs_uri)
        artifact_store = Client().active_stack.artifact_store
                    with artifact_store.open(self.logs_uri, "a") as file:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(visualization_uri, "w") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self.data_path, "rb") as file_:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self.data_path, "wb") as file_:
        artifact_store = Client().active_stack.artifact_store
        if not artifact_store.exists(
            self.data_path
        ) and not artifact_store.exists(self.metadata_path):
        if artifact_store.exists(self.data_path):
        artifact_store = Client().active_stack.artifact_store
                artifact_store.mkdir(element_path)
            if artifact_store.exists(self.metadata_path):
                artifact_store.remove(self.metadata_path)
                artifact_store.rmtree(entry["path"])
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "rb") as fid:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "wb") as fid:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        if artifact_store.exists(numpy_file):
            with artifact_store.open(numpy_file, "rb") as f:
        elif artifact_store.exists(os.path.join(self.uri, DATA_FILENAME)):
                with artifact_store.open(
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(
            os.path.join(self.uri, NUMPY_FILENAME), "wb"
        ) as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(output_path, "wb") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(output_path, "wb") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        if artifact_store.exists(self.parquet_path):
                with artifact_store.open(self.parquet_path, mode="rb") as f:
            with artifact_store.open(self.csv_path, mode="rb") as f:
        artifact_store = Client().active_stack.artifact_store
            with artifact_store.open(self.parquet_path, mode="wb") as f:
            with artifact_store.open(self.csv_path, mode="wb") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(describe_uri, mode="wb") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "r") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "w") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self._get_filepath(data_type), "r") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self._get_filepath(type(data)), "w") as f:
from zenml.client import Client
    artifact_store = stack.artifact_store
        if artifact_store.exists(artifact_uri):
        artifact_store.makedirs(artifact_uri)
    artifact_store = Client().active_stack.artifact_store
        if artifact_store.isdir(artifact_uri):
            artifact_store.rmtree(artifact_uri)
#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
#  or implied. See the License for the specific language governing
#  permissions and limitations under the License.

import os
from pathlib import Path

import pytest

from zenml.client import Client


def test_files_outside_of_artifact_store_are_not_reachable_by_it(
    clean_client: "Client",
):
    """Tests that no operations outside of bounds of artifact store could happen."""
    a_s = clean_client.active_stack.artifact_store

    outside_dir = Path(a_s.path) / ".."
    outside_file = str(outside_dir / "tmp.file")
    try:
        # create a file outside of artifact store
        with open(outside_file, "w") as f:
            f.write("test")
        # try to open it via artifact store interface
        with pytest.raises(FileNotFoundError):
            a_s.open(outside_file, "r")
        # try to copy it via artifact store interface
        with pytest.raises(FileNotFoundError):
            a_s.copyfile(outside_file, ".", "r")
    except Exception as e:
        raise e
    finally:
        os.remove(outside_file)

    inside_file = str(Path(a_s.path) / "tmp.file")
    try:
        # create a file inside of artifact store
        with open(inside_file, "w") as f:
            f.write("test")
        # try to open it via artifact store interface
        assert a_s.open(inside_file, "r").read() == "test"
        # try to copy it via artifact store interface
        inside_file2 = str(Path(a_s.path) / "tmp2.file")
        a_s.copyfile(inside_file, inside_file2, "r")
        # try to open it via artifact store interface
        assert open(inside_file2, "r").read() == "test"
        # try to copy it via artifact store interface, but with target outside of bounds
        with pytest.raises(FileNotFoundError):
            a_s.copyfile(inside_file, ".", "r")
    except Exception as e:
        raise e
    finally:
        os.remove(inside_file)
        os.remove(inside_file2)
    tmp_path, clean_client_with_run: "Client"
from zenml.client import Client
def model_metadata_dir(model_artifact, clean_client: "Client"):
    temp_dir = tempfile.mkdtemp(
        dir=clean_client.active_stack.artifact_store.path
    )
def numpy_file_uri(clean_client: "Client"):
    temp_dir = tempfile.mkdtemp(
        dir=clean_client.active_stack.artifact_store.path
    )
from zenml.client import Client
def test_container_materializer_for_custom_types(
    mocker, clean_client: "Client"
):
    with TemporaryDirectory(
        dir=clean_client.active_stack.artifact_store.path
    ) as artifact_uri:
    with TemporaryDirectory(
        dir=clean_client.active_stack.artifact_store.path
    ) as artifact_uri:
from zenml.client import Client
    artifact_store_uri = Client().active_stack.artifact_store.path
    with TemporaryDirectory(dir=artifact_store_uri) as artifact_uri: