genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/secure_loading.py

@@ -0,0 +1,221 @@
+"""
+Genesis-Flow Secure Model Loading
+
+This module provides secure alternatives to pickle-based model loading,
+addressing security vulnerabilities in model deserialization.
+"""
+
+import io
+import pickle
+import cloudpickle
+import logging
+import hashlib
+from typing import Any, Set, Optional, Union, Type
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Allowlist of safe classes that can be unpickled
+SAFE_PICKLE_CLASSES = {
+    # NumPy types
+    'numpy.ndarray',
+    'numpy.dtype',
+    'numpy.int32', 'numpy.int64', 'numpy.float32', 'numpy.float64',
+    'numpy.bool_', 'numpy.str_',
+    
+    # Pandas types
+    'pandas.core.frame.DataFrame',
+    'pandas.core.series.Series',
+    'pandas.core.index.Index',
+    'pandas.core.dtypes.dtypes.CategoricalDtype',
+    
+    # Scikit-learn estimators (core models only)
+    'sklearn.linear_model._base.LinearRegression',
+    'sklearn.linear_model._logistic.LogisticRegression',
+    'sklearn.ensemble._forest.RandomForestClassifier',
+    'sklearn.ensemble._forest.RandomForestRegressor',
+    'sklearn.tree._classes.DecisionTreeClassifier',
+    'sklearn.tree._classes.DecisionTreeRegressor',
+    'sklearn.svm._classes.SVC',
+    'sklearn.svm._classes.SVR',
+    
+    # Standard Python types
+    'builtins.dict', 'builtins.list', 'builtins.tuple', 'builtins.set',
+    'builtins.str', 'builtins.int', 'builtins.float', 'builtins.bool',
+    
+    # Collections
+    'collections.OrderedDict',
+    'collections.defaultdict',
+    
+    # Genesis-Flow internal types
+    'mlflow.models.signature.ModelSignature',
+    'mlflow.types.schema.Schema',
+}
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """
+    Secure unpickler that only allows safe, whitelisted classes.
+    
+    This prevents arbitrary code execution during model deserialization
+    by restricting which classes can be instantiated.
+    """
+    
+    def __init__(self, file, *, safe_classes: Optional[Set[str]] = None):
+        super().__init__(file)
+        self.safe_classes = safe_classes or SAFE_PICKLE_CLASSES
+        
+    def find_class(self, module: str, name: str) -> Type:
+        """
+        Override to restrict class loading to safe classes only.
+        
+        Args:
+            module: Module name
+            name: Class name
+            
+        Returns:
+            Class object if safe
+            
+        Raises:
+            SecurityError: If class is not in allowlist
+        """
+        full_name = f"{module}.{name}"
+        
+        # Check if the class is in our safe list
+        if full_name in self.safe_classes:
+            logger.debug(f"Loading safe class: {full_name}")
+            return super().find_class(module, name)
+        
+        # Additional check for known safe modules
+        safe_modules = {
+            'numpy': ['ndarray', 'dtype', 'int32', 'int64', 'float32', 'float64', 'bool_'],
+            'pandas': ['DataFrame', 'Series', 'Index'],
+            'builtins': ['dict', 'list', 'tuple', 'set', 'str', 'int', 'float', 'bool'],
+        }
+        
+        if module in safe_modules and name in safe_modules[module]:
+            logger.debug(f"Loading safe module class: {full_name}")
+            return super().find_class(module, name)
+        
+        # Log and block unsafe class
+        logger.warning(f"Blocked potentially unsafe class: {full_name}")
+        raise pickle.UnpicklingError(
+            f"Security: Class '{full_name}' is not in the allowlist. "
+            f"If this is a legitimate model class, add it to SAFE_PICKLE_CLASSES."
+        )
+
+class SecureModelLoader:
+    """
+    Secure model loading with multiple safety mechanisms.
+    """
+    
+    @staticmethod
+    def calculate_file_hash(file_path: Union[str, Path]) -> str:
+        """Calculate SHA256 hash of a file for integrity checking."""
+        hasher = hashlib.sha256()
+        with open(file_path, 'rb') as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hasher.update(chunk)
+        return hasher.hexdigest()
+    
+    @staticmethod
+    def safe_pickle_load(file_path: Union[str, Path], *, safe_classes: Optional[Set[str]] = None) -> Any:
+        """
+        Safely load a pickle file using restricted unpickler.
+        
+        Args:
+            file_path: Path to pickle file
+            safe_classes: Optional custom set of safe classes
+            
+        Returns:
+            Unpickled object
+            
+        Raises:
+            SecurityError: If file contains unsafe classes
+            FileNotFoundError: If file doesn't exist
+        """
+        file_path = Path(file_path)
+        
+        if not file_path.exists():
+            raise FileNotFoundError(f"Model file not found: {file_path}")
+        
+        # Log file hash for audit trail
+        file_hash = SecureModelLoader.calculate_file_hash(file_path)
+        logger.info(f"Loading model file {file_path.name} (SHA256: {file_hash[:16]}...)")
+        
+        try:
+            with open(file_path, 'rb') as f:
+                unpickler = RestrictedUnpickler(f, safe_classes=safe_classes)
+                model = unpickler.load()
+                logger.info(f"Successfully loaded model of type: {type(model).__name__}")
+                return model
+        except pickle.UnpicklingError as e:
+            logger.error(f"Security: Unsafe model file rejected: {e}")
+            raise SecurityError(f"Model file contains unsafe content: {e}") from e
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            raise
+    
+    @staticmethod
+    def safe_cloudpickle_load(file_path: Union[str, Path], *, safe_classes: Optional[Set[str]] = None) -> Any:
+        """
+        Safely load a cloudpickle file using restricted unpickler.
+        
+        Args:
+            file_path: Path to cloudpickle file
+            safe_classes: Optional custom set of safe classes
+            
+        Returns:
+            Unpickled object
+        """
+        file_path = Path(file_path)
+        
+        if not file_path.exists():
+            raise FileNotFoundError(f"Model file not found: {file_path}")
+        
+        # Calculate file hash for audit trail
+        file_hash = SecureModelLoader.calculate_file_hash(file_path)
+        logger.info(f"Loading cloudpickle file {file_path.name} (SHA256: {file_hash[:16]}...)")
+        
+        try:
+            with open(file_path, 'rb') as f:
+                # Use cloudpickle's load with our custom unpickler
+                # Note: This is a simplified approach - in practice, cloudpickle's
+                # load function would need to be modified to accept a custom unpickler
+                unpickler = RestrictedUnpickler(f, safe_classes=safe_classes)
+                model = unpickler.load()
+                logger.info(f"Successfully loaded cloudpickle model of type: {type(model).__name__}")
+                return model
+        except pickle.UnpicklingError as e:
+            logger.error(f"Security: Unsafe cloudpickle file rejected: {e}")
+            raise SecurityError(f"CloudPickle file contains unsafe content: {e}") from e
+        except Exception as e:
+            logger.error(f"Error loading cloudpickle model: {e}")
+            raise
+
+class SecurityError(Exception):
+    """Exception raised for security-related model loading issues."""
+    pass
+
+def add_safe_class(class_name: str) -> None:
+    """
+    Add a class to the safe loading allowlist.
+    
+    Args:
+        class_name: Full class name (e.g., 'mymodule.MyClass')
+    """
+    SAFE_PICKLE_CLASSES.add(class_name)
+    logger.info(f"Added {class_name} to safe loading allowlist")
+
+def remove_safe_class(class_name: str) -> None:
+    """
+    Remove a class from the safe loading allowlist.
+    
+    Args:
+        class_name: Full class name to remove
+    """
+    SAFE_PICKLE_CLASSES.discard(class_name)
+    logger.info(f"Removed {class_name} from safe loading allowlist")
+
+def get_safe_classes() -> Set[str]:
+    """Return a copy of the current safe classes set."""
+    return SAFE_PICKLE_CLASSES.copy()

mlflow/utils/security_validation.py

@@ -0,0 +1,384 @@
+"""
+Genesis-Flow Security Validation
+
+This module provides comprehensive input validation to prevent
+injection attacks, path traversal, and other security vulnerabilities.
+"""
+
+import os
+import re
+import logging
+from pathlib import Path
+from typing import Union, Optional
+from urllib.parse import urlparse
+
+logger = logging.getLogger(__name__)
+
+class SecurityValidationError(Exception):
+    """Exception raised for security validation failures."""
+    pass
+
+class InputValidator:
+    """
+    Comprehensive input validation for Genesis-Flow.
+    """
+    
+    # Safe characters for different input types
+    SAFE_NAME_PATTERN = re.compile(r'^[a-zA-Z0-9_\-\.]+$')
+    SAFE_TAG_KEY_PATTERN = re.compile(r'^[a-zA-Z0-9_\-\.\/]+$')
+    SAFE_TAG_VALUE_PATTERN = re.compile(r'^[a-zA-Z0-9_\-\.\s\/\:\@]+$')
+    
+    # Maximum lengths to prevent DoS
+    MAX_NAME_LENGTH = 256
+    MAX_TAG_KEY_LENGTH = 250
+    MAX_TAG_VALUE_LENGTH = 5000
+    MAX_DESCRIPTION_LENGTH = 10000
+    MAX_ARTIFACT_PATH_LENGTH = 1000
+    
+    @staticmethod
+    def validate_experiment_name(name: str) -> str:
+        """
+        Validate experiment name against injection attacks.
+        
+        Args:
+            name: Experiment name to validate
+            
+        Returns:
+            Validated name
+            
+        Raises:
+            SecurityValidationError: If name is invalid
+        """
+        if not name or not isinstance(name, str):
+            raise SecurityValidationError("Experiment name must be a non-empty string")
+        
+        if len(name) > InputValidator.MAX_NAME_LENGTH:
+            raise SecurityValidationError(
+                f"Experiment name too long (max {InputValidator.MAX_NAME_LENGTH} chars)"
+            )
+        
+        # Check for dangerous characters
+        if not InputValidator.SAFE_NAME_PATTERN.match(name):
+            raise SecurityValidationError(
+                "Experiment name contains invalid characters. "
+                "Only alphanumeric, dash, underscore, and dot allowed."
+            )
+        
+        # Prevent directory traversal
+        if '..' in name or '/' in name or '\\' in name:
+            raise SecurityValidationError("Experiment name cannot contain path traversal sequences")
+        
+        logger.debug(f"Validated experiment name: {name}")
+        return name
+    
+    @staticmethod
+    def validate_run_name(name: Optional[str]) -> Optional[str]:
+        """
+        Validate run name.
+        
+        Args:
+            name: Run name to validate (can be None)
+            
+        Returns:
+            Validated name or None
+        """
+        if name is None:
+            return None
+            
+        if not isinstance(name, str):
+            raise SecurityValidationError("Run name must be a string")
+        
+        if len(name) > InputValidator.MAX_NAME_LENGTH:
+            raise SecurityValidationError(
+                f"Run name too long (max {InputValidator.MAX_NAME_LENGTH} chars)"
+            )
+        
+        # Allow more flexible run names but still prevent injection
+        if '..' in name or '\x00' in name:
+            raise SecurityValidationError("Run name contains invalid sequences")
+        
+        return name
+    
+    @staticmethod
+    def validate_tag_key(key: str) -> str:
+        """
+        Validate tag key.
+        
+        Args:
+            key: Tag key to validate
+            
+        Returns:
+            Validated key
+        """
+        if not key or not isinstance(key, str):
+            raise SecurityValidationError("Tag key must be a non-empty string")
+        
+        if len(key) > InputValidator.MAX_TAG_KEY_LENGTH:
+            raise SecurityValidationError(
+                f"Tag key too long (max {InputValidator.MAX_TAG_KEY_LENGTH} chars)"
+            )
+        
+        if not InputValidator.SAFE_TAG_KEY_PATTERN.match(key):
+            raise SecurityValidationError(
+                "Tag key contains invalid characters"
+            )
+        
+        return key
+    
+    @staticmethod
+    def validate_tag_value(value: str) -> str:
+        """
+        Validate tag value.
+        
+        Args:
+            value: Tag value to validate
+            
+        Returns:
+            Validated value
+        """
+        if not isinstance(value, str):
+            raise SecurityValidationError("Tag value must be a string")
+        
+        if len(value) > InputValidator.MAX_TAG_VALUE_LENGTH:
+            raise SecurityValidationError(
+                f"Tag value too long (max {InputValidator.MAX_TAG_VALUE_LENGTH} chars)"
+            )
+        
+        # Check for control characters and potential injection
+        if '\x00' in value or '\r' in value:
+            raise SecurityValidationError("Tag value contains invalid control characters")
+        
+        return value
+    
+    @staticmethod
+    def validate_artifact_path(path: str, base_path: Optional[str] = None) -> str:
+        """
+        Validate artifact path to prevent path traversal attacks.
+        
+        Args:
+            path: Artifact path to validate
+            base_path: Optional base path for additional validation
+            
+        Returns:
+            Validated path
+            
+        Raises:
+            SecurityValidationError: If path is unsafe
+        """
+        if not path or not isinstance(path, str):
+            raise SecurityValidationError("Artifact path must be a non-empty string")
+        
+        if len(path) > InputValidator.MAX_ARTIFACT_PATH_LENGTH:
+            raise SecurityValidationError(
+                f"Artifact path too long (max {InputValidator.MAX_ARTIFACT_PATH_LENGTH} chars)"
+            )
+        
+        # Normalize the path to resolve any relative components
+        normalized_path = os.path.normpath(path)
+        
+        # Check for path traversal attempts
+        if '..' in normalized_path or normalized_path.startswith('/'):
+            raise SecurityValidationError("Artifact path contains path traversal sequences")
+        
+        # Additional check if base path provided
+        if base_path:
+            base = Path(base_path).resolve()
+            try:
+                target = (base / normalized_path).resolve()
+                # Ensure target is within base directory
+                target.relative_to(base)
+            except ValueError:
+                raise SecurityValidationError("Artifact path escapes base directory")
+        
+        logger.debug(f"Validated artifact path: {normalized_path}")
+        return normalized_path
+    
+    @staticmethod
+    def validate_metric_key(key: str) -> str:
+        """
+        Validate metric key.
+        
+        Args:
+            key: Metric key to validate
+            
+        Returns:
+            Validated key
+        """
+        if not key or not isinstance(key, str):
+            raise SecurityValidationError("Metric key must be a non-empty string")
+        
+        if len(key) > InputValidator.MAX_TAG_KEY_LENGTH:
+            raise SecurityValidationError(
+                f"Metric key too long (max {InputValidator.MAX_TAG_KEY_LENGTH} chars)"
+            )
+        
+        # Metric keys should be more restrictive
+        if not re.match(r'^[a-zA-Z0-9_\-\.]+$', key):
+            raise SecurityValidationError(
+                "Metric key contains invalid characters. "
+                "Only alphanumeric, dash, underscore, and dot allowed."
+            )
+        
+        return key
+    
+    @staticmethod
+    def validate_param_key(key: str) -> str:
+        """
+        Validate parameter key.
+        
+        Args:
+            key: Parameter key to validate
+            
+        Returns:
+            Validated key
+        """
+        if not key or not isinstance(key, str):
+            raise SecurityValidationError("Parameter key must be a non-empty string")
+        
+        if len(key) > InputValidator.MAX_TAG_KEY_LENGTH:
+            raise SecurityValidationError(
+                f"Parameter key too long (max {InputValidator.MAX_TAG_KEY_LENGTH} chars)"
+            )
+        
+        # Parameter keys should be restrictive
+        if not re.match(r'^[a-zA-Z0-9_\-\.]+$', key):
+            raise SecurityValidationError(
+                "Parameter key contains invalid characters. "
+                "Only alphanumeric, dash, underscore, and dot allowed."
+            )
+        
+        return key
+    
+    @staticmethod
+    def validate_param_value(value: Union[str, int, float]) -> str:
+        """
+        Validate parameter value.
+        
+        Args:
+            value: Parameter value to validate
+            
+        Returns:
+            Validated value as string
+        """
+        if value is None:
+            raise SecurityValidationError("Parameter value cannot be None")
+        
+        # Convert to string if needed
+        str_value = str(value)
+        
+        if len(str_value) > InputValidator.MAX_TAG_VALUE_LENGTH:
+            raise SecurityValidationError(
+                f"Parameter value too long (max {InputValidator.MAX_TAG_VALUE_LENGTH} chars)"
+            )
+        
+        # Check for control characters
+        if '\x00' in str_value or '\r' in str_value:
+            raise SecurityValidationError("Parameter value contains invalid control characters")
+        
+        return str_value
+    
+    @staticmethod
+    def validate_model_name(name: str) -> str:
+        """
+        Validate registered model name.
+        
+        Args:
+            name: Model name to validate
+            
+        Returns:
+            Validated name
+        """
+        if not name or not isinstance(name, str):
+            raise SecurityValidationError("Model name must be a non-empty string")
+        
+        if len(name) > InputValidator.MAX_NAME_LENGTH:
+            raise SecurityValidationError(
+                f"Model name too long (max {InputValidator.MAX_NAME_LENGTH} chars)"
+            )
+        
+        # Model names should be safe for file systems and URLs
+        if not re.match(r'^[a-zA-Z0-9_\-\.]+$', name):
+            raise SecurityValidationError(
+                "Model name contains invalid characters. "
+                "Only alphanumeric, dash, underscore, and dot allowed."
+            )
+        
+        # Prevent directory traversal
+        if '..' in name:
+            raise SecurityValidationError("Model name cannot contain '..'")
+        
+        return name
+    
+    @staticmethod
+    def validate_uri(uri: str, allowed_schemes: Optional[set] = None) -> str:
+        """
+        Validate URI for safety.
+        
+        Args:
+            uri: URI to validate
+            allowed_schemes: Set of allowed URI schemes
+            
+        Returns:
+            Validated URI
+        """
+        if not uri or not isinstance(uri, str):
+            raise SecurityValidationError("URI must be a non-empty string")
+        
+        if allowed_schemes is None:
+            allowed_schemes = {'http', 'https', 'file', 's3', 'gs', 'azure'}
+        
+        try:
+            parsed = urlparse(uri)
+        except Exception as e:
+            raise SecurityValidationError(f"Invalid URI format: {e}")
+        
+        if parsed.scheme and parsed.scheme not in allowed_schemes:
+            raise SecurityValidationError(
+                f"URI scheme '{parsed.scheme}' not allowed. "
+                f"Allowed schemes: {sorted(allowed_schemes)}"
+            )
+        
+        # Check for potential SSRF attacks
+        if parsed.hostname:
+            # Block private/local addresses for HTTP schemes
+            if parsed.scheme in ('http', 'https'):
+                if parsed.hostname.lower() in ('localhost', '127.0.0.1', '0.0.0.0'):
+                    raise SecurityValidationError("URI points to localhost/loopback address")
+                
+                # Block private IP ranges (basic check)
+                if (parsed.hostname.startswith('192.168.') or 
+                    parsed.hostname.startswith('10.') or
+                    parsed.hostname.startswith('172.')):
+                    raise SecurityValidationError("URI points to private IP address")
+        
+        return uri
+
+def secure_filename(filename: str) -> str:
+    """
+    Secure a filename by removing/replacing dangerous characters.
+    
+    Args:
+        filename: Original filename
+        
+    Returns:
+        Secured filename
+    """
+    if not filename:
+        raise SecurityValidationError("Filename cannot be empty")
+    
+    # Remove path components
+    filename = os.path.basename(filename)
+    
+    # Replace dangerous characters
+    filename = re.sub(r'[^\w\-_\.]', '_', filename)
+    
+    # Ensure it doesn't start with dot (hidden file)
+    if filename.startswith('.'):
+        filename = '_' + filename[1:]
+    
+    # Ensure reasonable length
+    if len(filename) > 255:
+        name, ext = os.path.splitext(filename)
+        filename = name[:251-len(ext)] + ext
+    
+    return filename

mlflow/utils/server_cli_utils.py

@@ -0,0 +1,61 @@
+"""
+Utilities for MLflow cli server config validation and resolving.
+NOTE: these functions are intended to be used as utilities for the cli click-based interface.
+Do not use for any other purpose as the potential Exceptions being raised will be misleading
+for users.
+"""
+
+import click
+
+from mlflow.store.tracking import DEFAULT_ARTIFACTS_URI, DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+from mlflow.utils.logging_utils import eprint
+from mlflow.utils.uri import is_local_uri
+
+
+def resolve_default_artifact_root(
+    serve_artifacts: bool,
+    default_artifact_root: str,
+    backend_store_uri: str,
+    resolve_to_local: bool = False,
+) -> str:
+    if serve_artifacts and not default_artifact_root:
+        default_artifact_root = DEFAULT_ARTIFACTS_URI
+    elif not serve_artifacts and not default_artifact_root:
+        if is_local_uri(backend_store_uri):
+            default_artifact_root = backend_store_uri
+        elif resolve_to_local:
+            default_artifact_root = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+        else:
+            msg = (
+                "Option 'default-artifact-root' is required when backend store is not "
+                "local file based and artifact serving is disabled."
+            )
+            eprint(msg)
+            raise click.UsageError(message=msg)
+    return default_artifact_root
+
+
+def _is_default_backend_store_uri(backend_store_uri: str) -> bool:
+    """Utility function to validate if the configured backend store uri location is set as the
+    default value for MLflow server.
+
+    Args:
+        backend_store_uri: The value set for the backend store uri for MLflow server artifact
+            handling.
+
+    Returns:
+        bool True if the default value is set.
+
+    """
+    return backend_store_uri == DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+
+
+def artifacts_only_config_validation(artifacts_only: bool, backend_store_uri: str) -> None:
+    if artifacts_only and not _is_default_backend_store_uri(backend_store_uri):
+        msg = (
+            "You are starting a tracking server in `--artifacts-only` mode and have provided a "
+            f"value for `--backend_store_uri`: '{backend_store_uri}'. A tracking server in "
+            "`--artifacts-only` mode cannot have a value set for `--backend_store_uri` to "
+            "properly proxy access to the artifact storage location."
+        )
+        raise click.UsageError(message=msg)

mlflow/utils/spark_utils.py

@@ -0,0 +1,15 @@
+def is_spark_connect_mode():
+    try:
+        from pyspark.sql.utils import is_remote
+    except ImportError:
+        return False
+    return is_remote()
+
+
+def get_spark_dataframe_type():
+    if is_spark_connect_mode():
+        from pyspark.sql.connect.dataframe import DataFrame as SparkDataFrame
+    else:
+        from pyspark.sql import DataFrame as SparkDataFrame
+
+    return SparkDataFrame