proxilion 0.0.1__py3-none-any.whl

proxilion/security/behavioral_drift.py

@@ -0,0 +1,794 @@
+"""
+Behavioral Drift Detection for Proxilion.
+
+Addresses OWASP ASI10: Rogue Agents.
+
+This module detects when an agent's behavior deviates significantly
+from its established baseline, potentially indicating:
+- Compromise or injection attack
+- Malfunction or loop behavior
+- Goal hijacking
+- Rogue agent behavior
+
+Example:
+    >>> from proxilion.security.behavioral_drift import (
+    ...     BehavioralMonitor,
+    ...     DriftDetector,
+    ...     KillSwitch,
+    ... )
+    >>>
+    >>> # Create monitor
+    >>> monitor = BehavioralMonitor(agent_id="my_agent")
+    >>>
+    >>> # Record normal behavior during baseline period
+    >>> for i in range(100):
+    ...     monitor.record_event("tool_call", {"tool": "search"})
+    ...     monitor.record_event("response", {"length": 150})
+    >>>
+    >>> # Lock baseline
+    >>> monitor.lock_baseline()
+    >>>
+    >>> # Detect drift during operation
+    >>> drift = monitor.check_drift()
+    >>> if drift.is_drifting:
+    ...     print(f"Behavioral drift detected: {drift.reason}")
+    ...     if drift.severity > 0.8:
+    ...         kill_switch.activate("Severe behavioral drift")
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import statistics
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Callable
+
+from proxilion.exceptions import BehavioralDriftError, EmergencyHaltError
+
+logger = logging.getLogger(__name__)
+
+
+class DriftMetric(Enum):
+    """Types of behavioral metrics tracked."""
+
+    TOOL_CALL_RATE = "tool_call_rate"
+    """Calls per minute."""
+
+    RESPONSE_LENGTH = "response_length"
+    """Average response length."""
+
+    ERROR_RATE = "error_rate"
+    """Errors per minute."""
+
+    UNIQUE_TOOLS = "unique_tools"
+    """Number of unique tools used."""
+
+    LATENCY = "latency"
+    """Average response latency."""
+
+    TOKEN_USAGE = "token_usage"
+    """Tokens consumed per request."""
+
+    TOOL_REPETITION = "tool_repetition"
+    """Same tool called consecutively."""
+
+    SCOPE_VIOLATIONS = "scope_violations"
+    """Attempts to exceed scope."""
+
+    CONTEXT_SIZE = "context_size"
+    """Size of conversation context."""
+
+    CUSTOM = "custom"
+    """User-defined metric."""
+
+
+@dataclass
+class MetricValue:
+    """A single metric measurement."""
+
+    metric: DriftMetric
+    value: float
+    timestamp: float
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class BaselineStats:
+    """Statistical baseline for a metric."""
+
+    metric: DriftMetric
+    mean: float
+    std_dev: float
+    min_value: float
+    max_value: float
+    sample_count: int
+    percentile_95: float
+    percentile_99: float
+
+    def z_score(self, value: float) -> float:
+        """Calculate z-score for a value."""
+        if self.std_dev == 0:
+            return 0.0 if value == self.mean else float("inf")
+        return (value - self.mean) / self.std_dev
+
+    def is_anomaly(self, value: float, threshold: float = 3.0) -> bool:
+        """Check if value is anomalous (beyond threshold std devs)."""
+        return abs(self.z_score(value)) > threshold
+
+
+@dataclass
+class DriftResult:
+    """Result of drift detection."""
+
+    is_drifting: bool
+    severity: float  # 0.0 to 1.0
+    drifting_metrics: list[tuple[DriftMetric, float, float]]  # (metric, value, z_score)
+    reason: str
+    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dict."""
+        return {
+            "is_drifting": self.is_drifting,
+            "severity": self.severity,
+            "drifting_metrics": [
+                {"metric": m.value, "value": v, "z_score": z}
+                for m, v, z in self.drifting_metrics
+            ],
+            "reason": self.reason,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+
+class BehavioralMonitor:
+    """
+    Monitors agent behavior and detects drift from baseline.
+
+    Tracks multiple behavioral metrics and uses statistical analysis
+    to detect when current behavior deviates from established patterns.
+
+    Example:
+        >>> monitor = BehavioralMonitor(agent_id="my_agent")
+        >>>
+        >>> # Record events during operation
+        >>> monitor.record_tool_call("search", {"query": "test"})
+        >>> monitor.record_response({"content": "result", "tokens": 50})
+        >>>
+        >>> # Check for drift
+        >>> result = monitor.check_drift()
+        >>> if result.is_drifting:
+        ...     handle_drift(result)
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        baseline_window: int = 100,
+        detection_window: int = 10,
+        drift_threshold: float = 3.0,
+        min_baseline_samples: int = 20,
+    ) -> None:
+        """
+        Initialize the monitor.
+
+        Args:
+            agent_id: Unique identifier for the agent.
+            baseline_window: Number of samples for baseline calculation.
+            detection_window: Recent samples for drift detection.
+            drift_threshold: Z-score threshold for drift detection.
+            min_baseline_samples: Minimum samples before baseline is valid.
+        """
+        self.agent_id = agent_id
+        self._baseline_window = baseline_window
+        self._detection_window = detection_window
+        self._drift_threshold = drift_threshold
+        self._min_baseline_samples = min_baseline_samples
+
+        # Metric storage
+        self._metrics: dict[DriftMetric, deque[MetricValue]] = {}
+        for metric in DriftMetric:
+            self._metrics[metric] = deque(maxlen=baseline_window)
+
+        # Baseline (locked after initial period)
+        self._baseline: dict[DriftMetric, BaselineStats] = {}
+        self._baseline_locked = False
+
+        # Rate tracking
+        self._event_times: deque[float] = deque(maxlen=1000)
+        self._tool_history: deque[str] = deque(maxlen=100)
+        self._error_count = 0
+
+        # Callbacks
+        self._drift_callbacks: list[Callable[[DriftResult], None]] = []
+
+        self._lock = threading.RLock()
+
+        logger.debug(f"BehavioralMonitor initialized for agent: {agent_id}")
+
+    def record_event(
+        self,
+        event_type: str,
+        data: dict[str, Any],
+    ) -> None:
+        """
+        Record a generic event.
+
+        Args:
+            event_type: Type of event (tool_call, response, error, etc.).
+            data: Event data.
+        """
+        now = time.time()
+        self._event_times.append(now)
+
+        with self._lock:
+            if event_type == "tool_call":
+                self._record_tool_call(data, now)
+            elif event_type == "response":
+                self._record_response(data, now)
+            elif event_type == "error":
+                self._record_error(data, now)
+            elif event_type == "latency":
+                self._record_metric(DriftMetric.LATENCY, data.get("value", 0), now)
+            elif event_type == "tokens":
+                self._record_metric(DriftMetric.TOKEN_USAGE, data.get("value", 0), now)
+            elif event_type == "context_size":
+                self._record_metric(DriftMetric.CONTEXT_SIZE, data.get("value", 0), now)
+            elif event_type == "scope_violation":
+                self._record_metric(DriftMetric.SCOPE_VIOLATIONS, 1.0, now)
+
+    def record_tool_call(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any] | None = None,
+        latency_ms: float | None = None,
+    ) -> None:
+        """Record a tool call event."""
+        now = time.time()
+        self._event_times.append(now)
+
+        with self._lock:
+            self._tool_history.append(tool_name)
+
+            # Calculate call rate (calls per minute)
+            recent_calls = sum(1 for t in self._event_times if now - t < 60)
+            self._record_metric(DriftMetric.TOOL_CALL_RATE, recent_calls, now)
+
+            # Track unique tools
+            unique_tools = len(set(self._tool_history))
+            self._record_metric(DriftMetric.UNIQUE_TOOLS, unique_tools, now)
+
+            # Track repetition
+            if len(self._tool_history) >= 2:
+                repetition = sum(
+                    1 for i in range(1, len(self._tool_history))
+                    if self._tool_history[i] == self._tool_history[i - 1]
+                )
+                self._record_metric(DriftMetric.TOOL_REPETITION, repetition, now)
+
+            # Record latency if provided
+            if latency_ms is not None:
+                self._record_metric(DriftMetric.LATENCY, latency_ms, now)
+
+    def record_response(
+        self,
+        response: dict[str, Any],
+    ) -> None:
+        """Record a response event."""
+        now = time.time()
+
+        with self._lock:
+            # Response length
+            content = response.get("content", "")
+            if isinstance(content, str):
+                self._record_metric(DriftMetric.RESPONSE_LENGTH, len(content), now)
+
+            # Token usage
+            tokens = response.get("tokens") or response.get("token_count")
+            if tokens:
+                self._record_metric(DriftMetric.TOKEN_USAGE, tokens, now)
+
+    def record_error(self, error_info: dict[str, Any]) -> None:
+        """Record an error event."""
+        now = time.time()
+        self._error_count += 1
+
+        with self._lock:
+            # Error rate (errors per minute)
+            recent_errors = sum(
+                1 for mv in self._metrics[DriftMetric.ERROR_RATE]
+                if now - mv.timestamp < 60
+            )
+            self._record_metric(DriftMetric.ERROR_RATE, recent_errors + 1, now)
+
+    def _record_tool_call(self, data: dict[str, Any], timestamp: float) -> None:
+        """Internal tool call recording."""
+        tool_name = data.get("tool") or data.get("tool_name", "unknown")
+        self._tool_history.append(tool_name)
+
+        # Calculate metrics
+        recent_calls = sum(1 for t in self._event_times if timestamp - t < 60)
+        self._record_metric(DriftMetric.TOOL_CALL_RATE, recent_calls, timestamp)
+
+        unique_tools = len(set(self._tool_history))
+        self._record_metric(DriftMetric.UNIQUE_TOOLS, unique_tools, timestamp)
+
+    def _record_response(self, data: dict[str, Any], timestamp: float) -> None:
+        """Internal response recording."""
+        content = data.get("content", "")
+        if isinstance(content, str):
+            self._record_metric(DriftMetric.RESPONSE_LENGTH, len(content), timestamp)
+
+        tokens = data.get("tokens", 0)
+        if tokens:
+            self._record_metric(DriftMetric.TOKEN_USAGE, tokens, timestamp)
+
+    def _record_error(self, data: dict[str, Any], timestamp: float) -> None:
+        """Internal error recording."""
+        self._error_count += 1
+        recent_errors = sum(
+            1 for mv in self._metrics[DriftMetric.ERROR_RATE]
+            if timestamp - mv.timestamp < 60
+        )
+        self._record_metric(DriftMetric.ERROR_RATE, recent_errors + 1, timestamp)
+
+    def _record_metric(
+        self,
+        metric: DriftMetric,
+        value: float,
+        timestamp: float,
+    ) -> None:
+        """Record a metric value."""
+        self._metrics[metric].append(MetricValue(
+            metric=metric,
+            value=value,
+            timestamp=timestamp,
+        ))
+
+    def lock_baseline(self) -> dict[DriftMetric, BaselineStats]:
+        """
+        Lock the current baseline.
+
+        Calculates statistical baselines from current data and
+        locks them for future drift detection.
+
+        Returns:
+            Dictionary of baseline stats per metric.
+
+        Raises:
+            ValueError: If not enough samples for baseline.
+        """
+        with self._lock:
+            self._baseline = {}
+
+            for metric, values in self._metrics.items():
+                if len(values) < self._min_baseline_samples:
+                    continue
+
+                samples = [v.value for v in values]
+
+                # Calculate statistics
+                mean = statistics.mean(samples)
+                std_dev = statistics.stdev(samples) if len(samples) > 1 else 0.0
+                sorted_samples = sorted(samples)
+                p95_idx = int(len(sorted_samples) * 0.95)
+                p99_idx = int(len(sorted_samples) * 0.99)
+
+                self._baseline[metric] = BaselineStats(
+                    metric=metric,
+                    mean=mean,
+                    std_dev=std_dev,
+                    min_value=min(samples),
+                    max_value=max(samples),
+                    sample_count=len(samples),
+                    percentile_95=sorted_samples[p95_idx] if p95_idx < len(sorted_samples) else max(samples),
+                    percentile_99=sorted_samples[p99_idx] if p99_idx < len(sorted_samples) else max(samples),
+                )
+
+            self._baseline_locked = True
+            logger.info(f"Baseline locked with {len(self._baseline)} metrics")
+
+            return self._baseline
+
+    def check_drift(self) -> DriftResult:
+        """
+        Check for behavioral drift from baseline.
+
+        Returns:
+            DriftResult indicating if drift was detected.
+        """
+        with self._lock:
+            if not self._baseline_locked:
+                # Auto-lock baseline if we have enough samples
+                has_enough = any(
+                    len(values) >= self._min_baseline_samples
+                    for values in self._metrics.values()
+                )
+                if has_enough:
+                    self.lock_baseline()
+                else:
+                    return DriftResult(
+                        is_drifting=False,
+                        severity=0.0,
+                        drifting_metrics=[],
+                        reason="Baseline not yet established",
+                    )
+
+            drifting_metrics: list[tuple[DriftMetric, float, float]] = []
+            max_severity = 0.0
+
+            for metric, baseline in self._baseline.items():
+                # Get recent values
+                recent = list(self._metrics[metric])[-self._detection_window:]
+                if not recent:
+                    continue
+
+                # Calculate current value (average of recent)
+                current_value = statistics.mean([v.value for v in recent])
+
+                # Calculate z-score
+                z_score = baseline.z_score(current_value)
+
+                # Check for drift
+                if abs(z_score) > self._drift_threshold:
+                    drifting_metrics.append((metric, current_value, z_score))
+
+                    # Calculate severity (normalized z-score)
+                    severity = min(1.0, abs(z_score) / (self._drift_threshold * 2))
+                    max_severity = max(max_severity, severity)
+
+            if drifting_metrics:
+                reasons = [
+                    f"{m.value}: {v:.2f} (z={z:.1f})"
+                    for m, v, z in drifting_metrics
+                ]
+                result = DriftResult(
+                    is_drifting=True,
+                    severity=max_severity,
+                    drifting_metrics=drifting_metrics,
+                    reason=f"Drift detected in: {', '.join(reasons)}",
+                )
+
+                # Notify callbacks
+                for callback in self._drift_callbacks:
+                    try:
+                        callback(result)
+                    except Exception as e:
+                        logger.error(f"Drift callback error: {e}")
+
+                return result
+
+            return DriftResult(
+                is_drifting=False,
+                severity=0.0,
+                drifting_metrics=[],
+                reason="Behavior within normal parameters",
+            )
+
+    def on_drift(self, callback: Callable[[DriftResult], None]) -> None:
+        """Register a callback for drift detection."""
+        self._drift_callbacks.append(callback)
+
+    def get_current_metrics(self) -> dict[str, float]:
+        """Get current metric values."""
+        with self._lock:
+            result = {}
+            for metric, values in self._metrics.items():
+                if values:
+                    recent = list(values)[-self._detection_window:]
+                    result[metric.value] = statistics.mean([v.value for v in recent])
+            return result
+
+    def get_baseline(self) -> dict[DriftMetric, BaselineStats]:
+        """Get the current baseline."""
+        with self._lock:
+            return self._baseline.copy()
+
+    def reset(self) -> None:
+        """Reset the monitor."""
+        with self._lock:
+            for values in self._metrics.values():
+                values.clear()
+            self._baseline = {}
+            self._baseline_locked = False
+            self._event_times.clear()
+            self._tool_history.clear()
+            self._error_count = 0
+
+
+class KillSwitch:
+    """
+    Emergency halt mechanism for rogue agent behavior.
+
+    Provides immediate shutdown capability when severe behavioral
+    drift or other anomalies are detected.
+
+    Example:
+        >>> kill_switch = KillSwitch()
+        >>>
+        >>> # Register halt handlers
+        >>> kill_switch.on_halt(lambda reason: cleanup_resources())
+        >>> kill_switch.on_halt(lambda reason: notify_operators(reason))
+        >>>
+        >>> # Activate when needed
+        >>> if drift.severity > 0.9:
+        ...     kill_switch.activate("Severe behavioral drift detected")
+    """
+
+    def __init__(
+        self,
+        auto_reset_seconds: float | None = None,
+    ) -> None:
+        """
+        Initialize the kill switch.
+
+        Args:
+            auto_reset_seconds: If set, auto-reset after this many seconds.
+        """
+        self._active = False
+        self._activation_time: datetime | None = None
+        self._activation_reason: str = ""
+        self._auto_reset_seconds = auto_reset_seconds
+
+        self._halt_callbacks: list[Callable[[str], None]] = []
+        self._reset_callbacks: list[Callable[[], None]] = []
+
+        self._lock = threading.RLock()
+
+        logger.debug("KillSwitch initialized")
+
+    @property
+    def is_active(self) -> bool:
+        """Check if kill switch is active."""
+        with self._lock:
+            if self._active and self._auto_reset_seconds:
+                # Check for auto-reset
+                if self._activation_time:
+                    elapsed = (datetime.now(timezone.utc) - self._activation_time).total_seconds()
+                    if elapsed > self._auto_reset_seconds:
+                        self._active = False
+                        self._activation_reason = ""
+                        logger.info("Kill switch auto-reset")
+            return self._active
+
+    @property
+    def reason(self) -> str:
+        """Get activation reason."""
+        return self._activation_reason
+
+    def activate(
+        self,
+        reason: str,
+        triggered_by: str = "system",
+        raise_exception: bool = True,
+    ) -> None:
+        """
+        Activate the kill switch.
+
+        Args:
+            reason: Why the kill switch was activated.
+            triggered_by: What triggered the activation.
+            raise_exception: If True, raise EmergencyHaltError.
+
+        Raises:
+            EmergencyHaltError: If raise_exception is True.
+        """
+        with self._lock:
+            self._active = True
+            self._activation_time = datetime.now(timezone.utc)
+            self._activation_reason = reason
+
+            logger.critical(
+                f"KILL SWITCH ACTIVATED: {reason} (triggered by: {triggered_by})"
+            )
+
+            # Notify handlers
+            for callback in self._halt_callbacks:
+                try:
+                    callback(reason)
+                except Exception as e:
+                    logger.error(f"Halt callback error: {e}")
+
+        if raise_exception:
+            raise EmergencyHaltError(reason=reason, triggered_by=triggered_by)
+
+    def reset(self) -> bool:
+        """
+        Reset the kill switch.
+
+        Returns:
+            True if was active and is now reset.
+        """
+        with self._lock:
+            was_active = self._active
+            self._active = False
+            self._activation_reason = ""
+            self._activation_time = None
+
+            if was_active:
+                logger.warning("Kill switch reset")
+                for callback in self._reset_callbacks:
+                    try:
+                        callback()
+                    except Exception as e:
+                        logger.error(f"Reset callback error: {e}")
+
+            return was_active
+
+    def check(self) -> None:
+        """
+        Check if kill switch is active and raise if so.
+
+        Raises:
+            EmergencyHaltError: If kill switch is active.
+        """
+        if self.is_active:
+            raise EmergencyHaltError(
+                reason=self._activation_reason,
+                triggered_by="kill_switch_check",
+            )
+
+    def on_halt(self, callback: Callable[[str], None]) -> None:
+        """Register a callback for when kill switch activates."""
+        self._halt_callbacks.append(callback)
+
+    def on_reset(self, callback: Callable[[], None]) -> None:
+        """Register a callback for when kill switch resets."""
+        self._reset_callbacks.append(callback)
+
+    def get_status(self) -> dict[str, Any]:
+        """Get kill switch status."""
+        with self._lock:
+            return {
+                "active": self._active,
+                "reason": self._activation_reason,
+                "activation_time": self._activation_time.isoformat() if self._activation_time else None,
+            }
+
+
+class DriftDetector:
+    """
+    High-level drift detector with integrated kill switch.
+
+    Combines behavioral monitoring with automatic response
+    to detected anomalies.
+
+    Example:
+        >>> detector = DriftDetector(
+        ...     agent_id="my_agent",
+        ...     auto_halt_threshold=0.9,
+        ... )
+        >>>
+        >>> # Record events
+        >>> detector.record_tool_call("search", {"query": "test"})
+        >>>
+        >>> # This will auto-halt if drift exceeds threshold
+        >>> detector.check()
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        auto_halt_threshold: float = 0.9,
+        warning_threshold: float = 0.5,
+        monitor_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Initialize the detector.
+
+        Args:
+            agent_id: Unique identifier for the agent.
+            auto_halt_threshold: Severity threshold for automatic halt.
+            warning_threshold: Severity threshold for warnings.
+            monitor_kwargs: Additional kwargs for BehavioralMonitor.
+        """
+        self.agent_id = agent_id
+        self._auto_halt_threshold = auto_halt_threshold
+        self._warning_threshold = warning_threshold
+
+        self._monitor = BehavioralMonitor(agent_id, **(monitor_kwargs or {}))
+        self._kill_switch = KillSwitch()
+
+        # Wire up automatic drift handling
+        self._monitor.on_drift(self._handle_drift)
+
+    @property
+    def monitor(self) -> BehavioralMonitor:
+        """Get the behavioral monitor."""
+        return self._monitor
+
+    @property
+    def kill_switch(self) -> KillSwitch:
+        """Get the kill switch."""
+        return self._kill_switch
+
+    def record_tool_call(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any] | None = None,
+        latency_ms: float | None = None,
+    ) -> None:
+        """Record a tool call and check for drift."""
+        self._kill_switch.check()  # Fail fast if halted
+        self._monitor.record_tool_call(tool_name, arguments, latency_ms)
+
+    def record_response(self, response: dict[str, Any]) -> None:
+        """Record a response and check for drift."""
+        self._kill_switch.check()
+        self._monitor.record_response(response)
+
+    def record_error(self, error_info: dict[str, Any]) -> None:
+        """Record an error and check for drift."""
+        self._monitor.record_error(error_info)
+
+    def record_event(self, event_type: str, data: dict[str, Any]) -> None:
+        """Record a generic event."""
+        self._kill_switch.check()
+        self._monitor.record_event(event_type, data)
+
+    def check(self) -> DriftResult:
+        """
+        Check for drift and respond accordingly.
+
+        Returns:
+            DriftResult from the check.
+
+        Raises:
+            EmergencyHaltError: If drift exceeds auto_halt_threshold.
+        """
+        self._kill_switch.check()
+        result = self._monitor.check_drift()
+
+        if result.is_drifting:
+            self._handle_drift(result)
+
+        return result
+
+    def _handle_drift(self, result: DriftResult) -> None:
+        """Handle detected drift."""
+        if result.severity >= self._auto_halt_threshold:
+            self._kill_switch.activate(
+                reason=f"Severe behavioral drift: {result.reason}",
+                triggered_by="drift_detector",
+            )
+        elif result.severity >= self._warning_threshold:
+            logger.warning(
+                f"Behavioral drift warning for {self.agent_id}: {result.reason}"
+            )
+
+    def lock_baseline(self) -> dict[DriftMetric, BaselineStats]:
+        """Lock the baseline."""
+        return self._monitor.lock_baseline()
+
+    def reset(self) -> None:
+        """Reset the detector."""
+        self._monitor.reset()
+        self._kill_switch.reset()
+
+    def get_status(self) -> dict[str, Any]:
+        """Get detector status."""
+        return {
+            "agent_id": self.agent_id,
+            "kill_switch": self._kill_switch.get_status(),
+            "current_metrics": self._monitor.get_current_metrics(),
+            "baseline_locked": self._monitor._baseline_locked,
+        }
+
+
+# Convenience exports
+__all__ = [
+    # Core classes
+    "BehavioralMonitor",
+    "DriftDetector",
+    "KillSwitch",
+    # Data classes
+    "DriftResult",
+    "BaselineStats",
+    "MetricValue",
+    "DriftMetric",
+]