models.py

import uuid
from typing import List, Dict, Any, Optional, Literal
from pydantic import BaseModel, Field


# ── Shared Cloud Resources ──
class CloudResource(BaseModel):
    id: str
    type: str
    status: str
    meta: Dict[str, Any] = {}


# ── Single-Agent Legacy Models ──
class SentinelAction(BaseModel):
    command: str
    target_id: str


class SentinelObservation(BaseModel):
    inventory: List[CloudResource]
    security_logs: List[str]
    terminal_output: str


class SentinelState(BaseModel):
    episode_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    step_count: int = 0
    current_task_id: str = "easy-lockdown"


# ── Multi-Agent Typed Actions (2026 Architecture) ──
class SentinelAgentAction(BaseModel):
    action_type: Literal[
        "scan_logs",
        "flag_threat",
        "share_intel",
        "block_ip",
        "close_port",
        "lockdown",
        "revoke_access",
        "restore_service",
    ]
    target_ip: Optional[str] = Field(None, description="IP address to block or scan")
    target_resource: Optional[str] = Field(None, description="Resource ID to act upon")
    intel_payload: Optional[str] = Field(None, description="Intel message from Scanner to Remediator")


class AttackerAgentAction(BaseModel):
    action_type: Literal[
        "port_scan",
        "privilege_escalation",
        "lateral_movement",
        "data_exfiltration",
        "hide_tracks",
    ]
    target_resource: str = Field(..., description="Cloud resource to attack")


# ── Per-Agent Observations (Partial Observability) ──
class ScannerObservation(BaseModel):
    ip_connections: List[str]       # Scanner sees IPs and logs
    flagged_threats: List[str]
    security_logs: List[str]
    shared_intel: List[str]
    terminal_output: str


class RemediatorObservation(BaseModel):
    open_ports: List[str]           # Remediator sees ports only
    locked_resources: List[str]
    received_intel: List[str]       # Intel shared by Scanner
    terminal_output: str


class AttackerObservation(BaseModel):
    visible_resources: List[CloudResource]  # Attacker sees unlocked resources
    current_position: str
    data_exfiltrated: bool
    terminal_output: str


# ── War Room State (Shared World) ──
class WarRoomState(BaseModel):
    episode_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    step_count: int = 0
    current_task_id: str = "red-vs-blue"
    is_compromised: bool = False
    data_exfiltrated: bool = False
    active_threats: List[str] = []
    blocked_ips: List[str] = []
    closed_ports: List[str] = []
    locked_resources: List[str] = []
    shared_intel: List[str] = []
    false_positives: int = 0
    uptime_score: float = 0.99
    mttr_steps: int = 0
    attacker_reward: float = 0.05
    blue_team_reward: float = 0.05


# ── Multi-Agent Step Input ──
class MultiAgentAction(BaseModel):
    attacker: AttackerAgentAction
    scanner: SentinelAgentAction
    remediator: SentinelAgentAction