backend/ai/models/attack_models.py

"""
Attack models for AegisLM Red Team Engine.

This module defines the core data structures for attacks,
attack types, and attack results using Pydantic for validation.
"""

from enum import Enum
from typing import List, Optional, Dict, Any
from datetime import datetime
from pydantic import BaseModel, Field


class AttackType(str, Enum):
    """Enumeration of supported attack types."""
    JAILBREAK = "jailbreak"
    PROMPT_INJECTION = "prompt_injection"
    TOXICITY_TRIGGER = "toxicity_trigger"
    HALLUCINATION_TRAP = "hallucination_trap"


class SeverityLevel(str, Enum):
    """Severity levels for attacks."""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class Attack(BaseModel):
    """Core attack model representing a single attack attempt."""
    id: str = Field(..., description="Unique identifier for the attack")
    attack_type: AttackType = Field(..., description="Type of attack")
    severity: SeverityLevel = Field(..., description="Severity level of the attack")
    prompt: str = Field(..., description="The attack prompt")
    context: Optional[str] = Field(None, description="Additional context for the attack")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
    generation_timestamp: datetime = Field(default_factory=datetime.now)
    iteration: int = Field(default=1, description="Iteration number for mutated attacks")
    parent_attack_id: Optional[str] = Field(None, description="ID of parent attack if mutated")


class AttackAttempt(BaseModel):
    """Represents a single attack attempt with its response."""
    attack: Attack = Field(..., description="The attack that was executed")
    response: str = Field(..., description="Model's response to the attack")
    response_timestamp: datetime = Field(default_factory=datetime.now)
    response_metadata: Dict[str, Any] = Field(default_factory=dict)


class AttackResult(BaseModel):
    """Result of an attack attempt including evaluation."""
    attempt: AttackAttempt = Field(..., description="The attack attempt")
    success: bool = Field(..., description="Whether the attack was successful")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence in the success determination")
    evaluation_details: Dict[str, Any] = Field(default_factory=dict, description="Detailed evaluation results")
    execution_time_ms: Optional[int] = Field(None, description="Execution time in milliseconds")


class AttackChain(BaseModel):
    """Chain of related attacks showing evolution."""
    chain_id: str = Field(..., description="Unique identifier for the attack chain")
    attacks: List[AttackResult] = Field(..., description="List of attacks in the chain")
    final_success: bool = Field(..., description="Whether any attack in the chain succeeded")
    successful_iteration: Optional[int] = Field(None, description="Iteration that succeeded")
    total_attempts: int = Field(..., description="Total number of attempts in the chain")