Beveiliging van dynamisch modelladen in productie

"""
Model loading security audit framework.
Tests model loading pipelines for integrity verification,
access control, and injection vulnerabilities.
"""
 
import hashlib
import json
import os
import time
from pathlib import Path
from dataclasses import dataclass
from typing import Optional, Any
from datetime import datetime, timezone
 
@dataclass
class ModelArtifact:
    """Represents a model artifact in the loading pipeline."""
    name: str
    version: str
    path: str
    format: str  # savedmodel, torchscript, onnx, pickle, safetensors
    size_bytes: int
    sha256: str
    signature: Optional[str] = None
    loaded_at: Optional[str] = None
 
class ModelLoadingAuditor:
    """
    Audit model loading infrastructure for security vulnerabilities.
    """
 
    # Model formats ranked by security risk
    FORMAT_RISK = {
        "pickle": "CRITICAL",   # Arbitrary code execution
        "pt": "CRITICAL",       # PyTorch (uses pickle)
        "pth": "CRITICAL",      # PyTorch checkpoint (uses pickle)
        "savedmodel": "HIGH",   # Can contain custom ops / py_function
        "mar": "HIGH",          # TorchServe archive with Python handlers
        "onnx": "MEDIUM",       # Generally safe but custom ops possible
        "torchscript": "MEDIUM",# Safer than pickle but not fully sandboxed
        "safetensors": "LOW",   # Designed for safe loading
        "tflite": "LOW",        # Restricted operation set
    }
 
    def __init__(self, model_dir: str):
        self.model_dir = Path(model_dir)
        self.findings: list[dict] = []
 
    def _add(self, severity: str, title: str, detail: str) -> None:
        self.findings.append({
            "severity": severity, "title": title, "detail": detail,
        })
 
    def audit_directory_permissions(self) -> None:
        """Check filesystem permissions on model directories."""
        if not self.model_dir.exists():
            self._add("ERROR", "Model directory not found",
                      f"{self.model_dir} does not exist")
            return
 
        # Check directory permissions
        dir_stat = self.model_dir.stat()
        mode = dir_stat.st_mode & 0o777
 
        if mode & 0o002:  # World-writable
            self._add(
                "CRITICAL",
                f"Model directory is world-writable: {self.model_dir}",
                "Any user on the system can inject malicious models. "
                "Model directory should be writable only by the deployment "
                "service account.",
            )
 
        if mode & 0o020:  # Group-writable
            self._add(
                "HIGH",
                f"Model directory is group-writable: {self.model_dir}",
                f"Group members can modify model artifacts. "
                f"GID: {dir_stat.st_gid}. Restrict to deployment account.",
            )
 
        # Check individual model files
        for model_file in self.model_dir.rglob("*"):
            if model_file.is_file():
                file_stat = model_file.stat()
                file_mode = file_stat.st_mode & 0o777
                if file_mode & 0o002:
                    self._add(
                        "CRITICAL",
                        f"World-writable model file: {model_file}",
                        "Any user can replace this model file with a "
                        "malicious version.",
                    )
 
    def audit_model_formats(self) -> None:
        """
        Identify model formats and their associated risks.
        Pickle-based formats allow arbitrary code execution on load.
        """
        for model_file in self.model_dir.rglob("*"):
            if not model_file.is_file():
                continue
 
            ext = model_file.suffix.lstrip(".")
            # Also check for pickle files without extension
            name = model_file.name
 
            risk = None
            if ext in self.FORMAT_RISK:
                risk = self.FORMAT_RISK[ext]
            elif name.endswith(".pkl") or name.endswith(".pickle"):
                risk = "CRITICAL"
 
            if risk and risk in ("CRITICAL", "HIGH"):
                self._add(
                    risk,
                    f"High-risk model format: {model_file.name} ({ext})",
                    f"File: {model_file}. "
                    f"{'Pickle-based formats execute arbitrary Python code during deserialization. ' if ext in ('pickle', 'pt', 'pth', 'pkl') else ''}"
                    f"Consider migrating to safetensors format for "
                    f"safe, zero-copy model loading.",
                )
 
    def audit_integrity_verification(self) -> None:
        """
        Check if model loading includes integrity verification.
        Look for hash manifests, signatures, or checksums.
        """
        # Look for integrity metadata files
        integrity_files = list(self.model_dir.rglob("**/checksums*")) + \
                         list(self.model_dir.rglob("**/manifest*")) + \
                         list(self.model_dir.rglob("**/*.sig")) + \
                         list(self.model_dir.rglob("**/DIGEST"))
 
        if not integrity_files:
            self._add(
                "HIGH",
                "No integrity verification metadata found",
                f"No checksum, manifest, or signature files found in "
                f"{self.model_dir}. Models can be replaced without "
                f"detection. Implement model artifact signing.",
            )
        else:
            for f in integrity_files:
                self._add(
                    "INFO",
                    f"Integrity file found: {f.name}",
                    f"Path: {f}. Verify that the serving framework "
                    f"actually validates this before loading.",
                )
 
    def audit_version_rollback_protection(self) -> None:
        """
        Check for rollback protection.
        Without it, an attacker can revert to a vulnerable or
        poisoned older model version.
        """
        model_dirs = sorted(
            [d for d in self.model_dir.iterdir() if d.is_dir()],
            key=lambda d: d.name,
        )
 
        if len(model_dirs) > 1:
            versions = []
            for d in model_dirs:
                try:
                    # Try to parse version from directory name
                    versions.append((d.name, d))
                except ValueError:
                    continue
 
            if len(versions) > 5:
                self._add(
                    "MEDIUM",
                    f"Many model versions retained: {len(versions)}",
                    "Retaining many old model versions increases the "
                    "attack surface for rollback attacks. Implement a "
                    "retention policy and minimum version enforcement.",
                )
 
    def scan_for_pickle_exploits(self, model_path: str) -> list[dict]:
        """
        Scan a pickle-based model file for known malicious patterns.
        Uses static analysis to detect common pickle exploitation.
        """
        findings = []
        suspicious_opcodes = {
            b"cos\n": "os module import (possible command execution)",
            b"csys\n": "sys module import",
            b"csubprocess\n": "subprocess module import (command execution)",
            b"cbuiltins\n": "builtins access",
            b"c__builtin__\n": "legacy builtins access",
            b"cposixpath\n": "filesystem path operations",
            b"csocket\n": "socket operations (network access)",
            b"curllib\n": "URL library (network access)",
        }
 
        try:
            with open(model_path, "rb") as f:
                # Read in chunks to handle large models
                content = f.read(10 * 1024 * 1024)  # First 10MB
 
            for pattern, description in suspicious_opcodes.items():
                if pattern in content:
                    findings.append({
                        "severity": "CRITICAL",
                        "title": f"Suspicious pickle opcode: {description}",
                        "detail": (
                            f"File {model_path} contains pickle opcodes "
                            f"that import {description}. This is a strong "
                            f"indicator of a malicious model."
                        ),
                    })
 
        except IOError as e:
            findings.append({
                "severity": "ERROR",
                "title": f"Cannot read model file: {model_path}",
                "detail": str(e),
            })
 
        return findings
 
    def run_audit(self) -> list[dict]:
        """Run complete model loading security audit."""
        self.findings = []
        self.audit_directory_permissions()
        self.audit_model_formats()
        self.audit_integrity_verification()
        self.audit_version_rollback_protection()
 
        # Scan pickle-based files
        for model_file in self.model_dir.rglob("*"):
            if model_file.suffix in (".pt", ".pth", ".pkl", ".pickle"):
                pickle_findings = self.scan_for_pickle_exploits(
                    str(model_file)
                )
                self.findings.extend(pickle_findings)
 
        return self.findings
 
if __name__ == "__main__":
    import sys
 
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "/models"
    auditor = ModelLoadingAuditor(model_dir)
    findings = auditor.run_audit()
    for f in findings:
        print(f"[{f['severity']}] {f['title']}")
        print(f"  {f['detail']}\n")

"""
Model integrity verification middleware for serving frameworks.
Verifies model artifacts against a signed manifest before allowing
the serving framework to load them.
"""
 
import hashlib
import json
import hmac
from pathlib import Path
from typing import Optional
from datetime import datetime, timezone
 
class ModelIntegrityVerifier:
    """
    Verify model artifact integrity before hot-loading.
    Designed as middleware between the model source and
    the serving framework's load operation.
    """
 
    def __init__(
        self,
        manifest_path: str,
        signing_key: bytes,
        min_version: Optional[int] = None,
    ):
        """
        Args:
            manifest_path: Path to the signed model manifest
            signing_key: HMAC signing key for manifest verification
            min_version: Minimum allowed model version (rollback protection)
        """
        self.manifest_path = Path(manifest_path)
        self.signing_key = signing_key
        self.min_version = min_version
        self._load_manifest()
 
    def _load_manifest(self) -> None:
        """Load and verify the model manifest."""
        content = self.manifest_path.read_text()
 
        # Manifest format: JSON body + HMAC signature on last line
        lines = content.strip().split("\n")
        signature_line = lines[-1]
        body = "\n".join(lines[:-1])
 
        if not signature_line.startswith("SIGNATURE:"):
            raise ValueError("Manifest missing HMAC signature")
 
        expected_sig = signature_line.split(":", 1)[1].strip()
        actual_sig = hmac.new(
            self.signing_key, body.encode(), hashlib.sha256,
        ).hexdigest()
 
        if not hmac.compare_digest(actual_sig, expected_sig):
            raise SecurityError(
                "Manifest signature verification failed. "
                "The manifest may have been tampered with."
            )
 
        self.manifest = json.loads(body)
 
    def verify_model(
        self,
        model_name: str,
        model_version: str,
        model_path: str,
    ) -> tuple[bool, str]:
        """
        Verify a model artifact against the signed manifest.
 
        Returns:
            (is_valid, reason)
        """
        # Check rollback protection
        if self.min_version is not None:
            try:
                version_num = int(model_version)
                if version_num < self.min_version:
                    return False, (
                        f"Version {model_version} is below minimum "
                        f"allowed version {self.min_version}. "
                        f"Rollback protection triggered."
                    )
            except ValueError:
                pass  # Non-numeric version, skip check
 
        # Find model in manifest
        key = f"{model_name}/{model_version}"
        expected = self.manifest.get("models", {}).get(key)
 
        if expected is None:
            return False, (
                f"Model {key} not found in signed manifest. "
                f"Unapproved model version — loading blocked."
            )
 
        # Verify file hashes
        model_base = Path(model_path)
        for file_entry in expected.get("files", []):
            rel_path = file_entry["path"]
            expected_hash = file_entry["sha256"]
            expected_size = file_entry["size"]
 
            file_path = model_base / rel_path
            if not file_path.exists():
                return False, f"Expected file missing: {rel_path}"
 
            # Verify size first (fast check)
            actual_size = file_path.stat().st_size
            if actual_size != expected_size:
                return False, (
                    f"Size mismatch for {rel_path}: "
                    f"expected {expected_size}, got {actual_size}"
                )
 
            # Verify hash
            h = hashlib.sha256()
            with open(file_path, "rb") as f:
                while True:
                    chunk = f.read(8192)
                    if not chunk:
                        break
                    h.update(chunk)
            actual_hash = h.hexdigest()
 
            if actual_hash != expected_hash:
                return False, (
                    f"Hash mismatch for {rel_path}: "
                    f"expected {expected_hash[:16]}..., "
                    f"got {actual_hash[:16]}..."
                )
 
        # Check model format safety
        model_format = expected.get("format", "unknown")
        if model_format in ("pickle", "pt", "pth"):
            return False, (
                f"Unsafe model format: {model_format}. "
                f"Only safetensors and ONNX are allowed in production."
            )
 
        return True, "Model artifact verified successfully"
 
    def create_signed_manifest(
        self,
        models_dir: str,
        output_path: str,
    ) -> str:
        """
        Create a signed manifest for all models in a directory.
        This should run in a trusted build environment.
        """
        manifest = {
            "created_at": datetime.now(timezone.utc).isoformat(),
            "models": {},
        }
 
        models_base = Path(models_dir)
        for model_dir in sorted(models_base.iterdir()):
            if not model_dir.is_dir():
                continue
            model_name = model_dir.name
 
            for version_dir in sorted(model_dir.iterdir()):
                if not version_dir.is_dir():
                    continue
                version = version_dir.name
                key = f"{model_name}/{version}"
 
                files = []
                for file_path in sorted(version_dir.rglob("*")):
                    if not file_path.is_file():
                        continue
 
                    h = hashlib.sha256()
                    with open(file_path, "rb") as f:
                        while True:
                            chunk = f.read(8192)
                            if not chunk:
                                break
                            h.update(chunk)
 
                    files.append({
                        "path": str(file_path.relative_to(version_dir)),
                        "sha256": h.hexdigest(),
                        "size": file_path.stat().st_size,
                    })
 
                manifest["models"][key] = {
                    "files": files,
                    "format": self._detect_format(version_dir),
                }
 
        body = json.dumps(manifest, indent=2)
        signature = hmac.new(
            self.signing_key, body.encode(), hashlib.sha256,
        ).hexdigest()
 
        output = f"{body}\nSIGNATURE:{signature}"
        Path(output_path).write_text(output)
        return output_path
 
    def _detect_format(self, model_dir: Path) -> str:
        """Detect model format from file extensions."""
        extensions = {f.suffix for f in model_dir.rglob("*") if f.is_file()}
        if ".safetensors" in extensions:
            return "safetensors"
        if ".onnx" in extensions:
            return "onnx"
        if ".pt" in extensions or ".pth" in extensions:
            return "pt"
        if ".pkl" in extensions:
            return "pickle"
        if "saved_model.pb" in {f.name for f in model_dir.rglob("*")}:
            return "savedmodel"
        return "unknown"

#!/usr/bin/env bash
# Monitor model loading events in production for security anomalies
 
MODEL_DIR="${1:-/models}"
LOG_FILE="/var/log/model-loading-audit.log"
 
echo "=== Model Loading Monitor ==="
echo "Watching: $MODEL_DIR"
echo "Log: $LOG_FILE"
 
# Use inotifywait to monitor filesystem changes
# Requires inotify-tools package
if ! command -v inotifywait &>/dev/null; then
    echo "Installing inotify-tools..."
    apt-get install -y inotify-tools 2>/dev/null || \
        yum install -y inotify-tools 2>/dev/null || \
        echo "Please install inotify-tools"
fi
 
log_event() {
    local severity="$1"
    local message="$2"
    local timestamp
    timestamp=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
    echo "${timestamp} [${severity}] ${message}" | tee -a "$LOG_FILE"
}
 
# Monitor for file changes in model directory
inotifywait -m -r -e create,modify,delete,moved_to \
    --format '%T %w%f %e' --timefmt '%Y-%m-%dT%H:%M:%S' \
    "$MODEL_DIR" 2>/dev/null | while read -r timestamp filepath event; do
 
    # Classify the event
    case "$event" in
        *CREATE*|*MOVED_TO*)
            if echo "$filepath" | grep -qE '\.(pt|pth|pkl|pickle|mar)$'; then
                log_event "CRITICAL" "High-risk model format created: $filepath"
            elif echo "$filepath" | grep -qE '\.(safetensors|onnx)$'; then
                log_event "INFO" "Model file created: $filepath"
            else
                log_event "LOW" "File created in model dir: $filepath"
            fi
            ;;
        *MODIFY*)
            log_event "HIGH" "Model file modified: $filepath"
            # Compute hash for audit trail
            if [ -f "$filepath" ]; then
                hash=$(sha256sum "$filepath" | awk '{print $1}')
                log_event "INFO" "New hash: $hash"
            fi
            ;;
        *DELETE*)
            log_event "MEDIUM" "Model file deleted: $filepath"
            ;;
    esac
done