avi2132/cyber-security-shield

1# models.py — Defines the 3 data structures used to communicate between the AI agent and the environment.2# Observation = what the agent sees, Action = what the agent does, Reward = feedback on how it did.3 4from pydantic import BaseModel, Field  # Pydantic gives us auto-validation and JSON serialization5from typing import Dict, Any, List, Optional, Union6 7 8class Observation(BaseModel):9    """What the environment shows the agent each step (current state snapshot)."""10 11    # Human-readable description of what's happening (e.g., "Incoming connection from IP: 192.168.1.42")12    text: str = Field(..., description="Textual description of the current state or command output.")13 14    # Machine-readable data the agent logic can parse (contents change per task)15    # e.g., {"ip": "192.168.1.42", "task": "ddos_sentinel"} or {"log": {...}, "task": "log_triager"}16    data: Dict[str, Any] = Field(default_factory=dict, description="Structured data (e.g., packet payloads, system state).")17 18    # Which tools the agent can use right now (e.g., ["block_ip", "rate_limit", "submit"])19    available_tools: List[str] = Field(default_factory=list, description="List of tools available for the current task.")20 21 22class Action(BaseModel):23    """What the agent decides to do — picks a tool and provides arguments."""24 25    # Tool to run. Options: "block_ip", "rate_limit", "analyze_header", "classify", "chmod", "ls", "cat", "submit", "check_domain_reputation", "quarantine_email", "allow_email", "get_users", "get_login_history", "disable_user"26    tool_name: str = Field(..., description="Name of the tool to execute.")27 28    # Arguments for that tool, e.g., {"ip": "192.168.1.42"} or {"classification": "rce_attempt"}29    tool_args: Dict[str, Any] = Field(default_factory=dict, description="Arguments for the executed tool.")30 31 32class Reward(BaseModel):33    """Feedback after each action — score + whether the episode ended + metrics."""34 35    # Step reward: positive = good action, negative = bad action, zero = neutral36    value: float = Field(default=0.0, description="Incremental or final reward signal.")37 38    # True when the episode is over (max steps reached, or agent submitted, or all tasks done)39    is_terminal: bool = Field(default=False, description="True if the environment reached a terminal state.")40 41    # Detailed metrics (fully populated only on the final/terminal step)42    # Includes: final_grade, fpr_percentage, success_rate, avg_latency_ms, etc.43    info: Dict[str, Any] = Field(default_factory=dict, description="Metrics: FPR, Uptime, Latency, Precision, etc.")44 45 46class State(BaseModel):47    """Episode metadata and environment state snapshot."""48 49    # Unique identifier for this episode50    episode_id: str = Field(..., description="Unique episode identifier (timestamp-based).")51 52    # Progress tracking53    steps_taken: int = Field(default=0, description="Number of steps executed in this episode.")54    max_steps: int = Field(default=50, description="Maximum allowed steps for this episode.")55    tasks_completed: int = Field(default=0, description="Number of tasks completed so far.")56    total_tasks: int = Field(default=5, description="Total number of tasks in this episode.")57    current_task: str = Field(default="", description="Name of the currently active task.")58 59    # Metrics60    avg_latency_ms: float = Field(default=0.0, description="Average latency per step (milliseconds).")61    step_latencies: List[float] = Field(default_factory=list, description="Latency of each individual step.")62 63    # Combined metrics from all tasks64    current_metrics: Dict[str, Any] = Field(default_factory=dict, description="Current metrics aggregated from all tasks.")65