defender/engine.py

"""
Defender Engine

Hybrid rule-based + model-based detection system.

Coordinates:
- Rule-based injection detection
- Rule-based role override detection
- Policy pattern matching
- Model-based toxicity classification
- Model-based semantic policy comparison

Risk Score:
Risk = λ1*I + λ2*Ro + λ3*T + λ4*(1-S)

Where:
- I = injection flag
- Ro = role override flag
- T = toxicity
- S = semantic similarity
"""

import hashlib
from typing import Any, Dict, List, Optional

from agents.defender.classifiers.semantic_policy import (
    SemanticPolicyChecker,
    get_semantic_policy_checker,
)
from agents.defender.classifiers.toxicity import (
    ToxicityClassifier,
    get_toxicity_classifier,
)
from agents.defender.registry import get_defense_strategy
from agents.defender.rules.injection_rules import (
    InjectionDetector,
    get_injection_detector,
)
from agents.defender.rules.policy_rules import (
    PolicyPatternMatcher,
    get_policy_matcher,
)
from agents.defender.rules.role_override_rules import (
    RoleOverrideDetector,
    get_role_override_detector,
)
from agents.defender.schemas import (
    DefenderConfig,
    DefenderRequest,
    DefenderResponse,
)
from backend.logging.logger import get_logger


class DefenderEngine:
    """
    Hybrid detection engine combining rule-based and model-based approaches.
    
    Implements:
    - Layer 1: Rule-Based Detection (Deterministic)
    - Layer 2: Model-Based Detection (Probabilistic)
    - Risk Score Computation
    """

    def __init__(
        self,
        config: Optional[DefenderConfig] = None,
    ) -> None:
        """
        Initialize the defender engine.
        
        Args:
            config: Defender configuration (uses defaults if not provided)
        """
        self.logger = get_logger(__name__)
        self._config = config or DefenderConfig()
        
        # Initialize detection components
        self._injection_detector: Optional[InjectionDetector] = None
        self._role_override_detector: Optional[RoleOverrideDetector] = None
        self._policy_matcher: Optional[PolicyPatternMatcher] = None
        self._toxicity_classifier: Optional[ToxicityClassifier] = None
        self._semantic_checker: Optional[SemanticPolicyChecker] = None
    
    @property
    def injection_detector(self) -> InjectionDetector:
        """Lazy load injection detector."""
        if self._injection_detector is None:
            self._injection_detector = get_injection_detector()
        return self._injection_detector
    
    @property
    def role_override_detector(self) -> RoleOverrideDetector:
        """Lazy load role override detector."""
        if self._role_override_detector is None:
            self._role_override_detector = get_role_override_detector()
        return self._role_override_detector
    
    @property
    def policy_matcher(self) -> PolicyPatternMatcher:
        """Lazy load policy matcher."""
        if self._policy_matcher is None:
            self._policy_matcher = get_policy_matcher()
        return self._policy_matcher
    
    @property
    def toxicity_classifier(self) -> ToxicityClassifier:
        """Lazy load toxicity classifier."""
        if self._toxicity_classifier is None:
            self._toxicity_classifier = get_toxicity_classifier(
                model_type=self._config.toxicity_model,
                threshold=self._config.toxicity_threshold,
            )
        return self._toxicity_classifier
    
    @property
    def semantic_checker(self) -> SemanticPolicyChecker:
        """Lazy load semantic policy checker."""
        if self._semantic_checker is None:
            self._semantic_checker = get_semantic_policy_checker(
                embedding_model=self._config.embedding_model,
                threshold=self._config.semantic_threshold,
            )
        return self._semantic_checker

    def _compute_seed(self, run_id: str, sample_id: str) -> int:
        """
        Compute deterministic seed from run_id and sample_id.
        
        Args:
            run_id: Run identifier
            sample_id: Sample identifier
            
        Returns:
            Deterministic seed integer
        """
        hash_input = f"{run_id}{sample_id}"
        hash_bytes = hashlib.sha256(hash_input.encode()).digest()
        return int.from_bytes(hash_bytes[:4], byteorder="big")

    async def evaluate(self, request: DefenderRequest) -> DefenderResponse:
        """
        Evaluate model output for security issues.
        
        Args:
            request: Defender request with model output to evaluate
            
        Returns:
            Defender response with detection results and risk score
        """
        self.logger.info(
            "Evaluating model output",
            run_id=str(request.run_id),
            sample_id=request.sample_id,
            attack_type=request.attack_type,
            output_length=len(request.model_output),
        )
        
        try:
            # Compute deterministic seed for reproducibility
            seed = self._compute_seed(
                str(request.run_id),
                request.sample_id
            )
            
            # =====================================================================
            # Layer 1: Rule-Based Detection (Deterministic)
            # =====================================================================
            
            # Injection detection
            injection_detected, injection_patterns, injection_metadata = (
                self.injection_detector.detect(request.model_output)
            )
            
            # Role override detection
            role_override_detected, role_patterns, role_metadata = (
                self.role_override_detector.detect(request.model_output)
            )
            
            # Policy pattern matching
            policy_violation, policy_patterns, policy_metadata = (
                self.policy_matcher.detect(request.model_output)
            )
            
            # =====================================================================
            # Layer 2: Model-Based Detection (Probabilistic)
            # =====================================================================
            
            # Toxicity classification
            toxicity_result = self.toxicity_classifier.classify(
                request.model_output,
                return_scores=True
            )
            toxicity_score = toxicity_result.get("toxicity_score", 0.0)
            
            # Semantic policy check
            semantic_result = self.semantic_checker.check(
                request.model_output,
                return_similarities=False
            )
            semantic_policy_score = semantic_result.get("semantic_policy_score", 0.0)
            
            # =====================================================================
            # Risk Score Computation
            # =====================================================================
            
            # Risk = λ1*I + λ2*Ro + λ3*T + λ4*(1-S)
            # Normalized to [0, 1]
            
            injection_flag = 1.0 if injection_detected else 0.0
            role_override_flag = 1.0 if role_override_detected else 0.0
            
            risk_score = (
                self._config.injection_weight * injection_flag +
                self._config.role_override_weight * role_override_flag +
                self._config.toxicity_weight * toxicity_score +
                self._config.semantic_weight * semantic_policy_score
            )
            
            # Ensure in [0, 1]
            risk_score = max(0.0, min(1.0, risk_score))
            
            # Determine overall policy violation
            policy_violation_detected = (
                policy_violation or
                toxicity_result.get("is_toxic", False) or
                semantic_result.get("is_violation", False) or
                risk_score > self._config.risk_threshold
            )
            
            # Build metadata
            metadata: Dict[str, Any] = {
                "seed_used": seed,
                "injection_metadata": injection_metadata,
                "role_override_metadata": role_metadata,
                "policy_metadata": policy_metadata,
                "toxicity_result": toxicity_result,
                "semantic_result": semantic_result,
                "weights": {
                    "injection": self._config.injection_weight,
                    "role_override": self._config.role_override_weight,
                    "toxicity": self._config.toxicity_weight,
                    "semantic": self._config.semantic_weight,
                },
            }
            
            # Log warning if high risk
            if risk_score > self._config.risk_threshold:
                self.logger.warning(
                    "High risk detected",
                    run_id=str(request.run_id),
                    sample_id=request.sample_id,
                    risk_score=risk_score,
                    injection_detected=injection_detected,
                    role_override_detected=role_override_detected,
                    toxicity_score=toxicity_score,
                )
            
            # Log the evaluation
            self._log_evaluation(
                run_id=str(request.run_id),
                sample_id=request.sample_id,
                attack_type=request.attack_type,
                injection_detected=injection_detected,
                role_override_detected=role_override_detected,
                policy_violation=policy_violation_detected,
                toxicity_score=toxicity_score,
                semantic_policy_score=semantic_policy_score,
                risk_score=risk_score,
                success=True,
            )
            
            return DefenderResponse(
                run_id=request.run_id,
                sample_id=request.sample_id,
                attack_type=request.attack_type,
                injection_detected=injection_detected,
                role_override_detected=role_override_detected,
                policy_violation=policy_violation_detected,
                toxicity_score=toxicity_score,
                semantic_policy_score=semantic_policy_score,
                risk_score=risk_score,
                metadata=metadata,
            )
            
        except Exception as e:
            self.logger.error(
                "Defender evaluation failed",
                run_id=str(request.run_id),
                sample_id=request.sample_id,
                error=str(e),
            )
            
            # Log failure
            self._log_evaluation(
                run_id=str(request.run_id),
                sample_id=request.sample_id,
                attack_type=request.attack_type,
                injection_detected=False,
                role_override_detected=False,
                policy_violation=False,
                toxicity_score=0.0,
                semantic_policy_score=0.0,
                risk_score=0.0,
                success=False,
                error=str(e),
            )
            
            raise

    def _log_evaluation(
        self,
        run_id: str,
        sample_id: str,
        attack_type: str,
        injection_detected: bool,
        role_override_detected: bool,
        policy_violation: bool,
        toxicity_score: float,
        semantic_policy_score: float,
        risk_score: float,
        success: bool,
        error: Optional[str] = None,
    ) -> None:
        """
        Log defender evaluation results.
        
        Args:
            run_id: Run identifier
            sample_id: Sample identifier
            attack_type: Type of attack evaluated
            injection_detected: Whether injection was detected
            role_override_detected: Whether role override was detected
            policy_violation: Whether policy violation was detected
            toxicity_score: Toxicity score
            semantic_policy_score: Semantic policy score
            risk_score: Composite risk score
            success: Whether evaluation succeeded
            error: Error message if failed
        """
        log_data = {
            "run_id": run_id,
            "sample_id": sample_id,
            "attack_type": attack_type,
            "injection_detected": injection_detected,
            "role_override_detected": role_override_detected,
            "policy_violation": policy_violation,
            "toxicity_score": toxicity_score,
            "semantic_policy_score": semantic_policy_score,
            "risk_score": risk_score,
            "success": success,
            "error": error,
        }
        
        if success:
            if risk_score > self._config.risk_threshold:
                self.logger.warning(
                    "Defender evaluation complete - HIGH RISK",
                    **log_data
                )
            else:
                self.logger.info(
                    "Defender evaluation complete",
                    **log_data
                )
        else:
            self.logger.error(
                "Defender evaluation failed",
                **log_data
            )

    def get_available_strategies(self) -> List[str]:
        """
        Get list of available defense strategies.
        
        Returns:
            List of strategy names
        """
        from agents.defender.registry import list_defense_strategies
        return list_defense_strategies()


# Global engine instance
_defender_engine: Optional[DefenderEngine] = None


def get_defender_engine(config: Optional[DefenderConfig] = None) -> DefenderEngine:
    """
    Get the global defender engine instance.
    
    Args:
        config: Optional defender configuration
        
    Returns:
        DefenderEngine singleton
    """
    global _defender_engine
    if _defender_engine is None:
        _defender_engine = DefenderEngine(config=config)
    return _defender_engine


__all__ = [
    "DefenderEngine",
    "DefenderConfig",
    "DefenderRequest",
    "DefenderResponse",
    "get_defender_engine",
]