bothari01/secops-env

1"""2SecOps Environment - Typed Models3 4Pydantic models for actions, observations, and state management.5"""6 7from typing import Any, Dict, List, Optional8from enum import Enum9from pydantic import BaseModel, Field, ConfigDict10 11 12class TaskType(str, Enum):13    """Available security operations tasks."""14 15    PII_REDACTION = "pii_redaction"16    PUBLIC_ACCESS = "public_access"17    GHOST_USER = "ghost_user"18    LOG_ANALYSIS = "log_analysis"19    CONFIG_HARDENING = "config_hardening"20 21 22class TaskDifficulty(str, Enum):23    """Task difficulty levels."""24 25    EASY = "easy"26    MEDIUM = "medium"27    HARD = "hard"28 29 30class ActionType(str, Enum):31    """Available action types."""32 33    ANALYZE = "analyze"34    IDENTIFY = "identify"35    APPLY_FIX = "apply_fix"36    FINALIZE = "finalize"37    NOOP = "noop"38    CLASSIFY = "classify"39    PRIORITIZE = "prioritize"40    REVIEW = "review"41    IDENTIFY_ISSUES = "identify_issues"42    SUGGEST_FIXES = "suggest_fixes"43    APPLY_FIXES = "apply_fixes"44 45 46class SecOpsAction(BaseModel):47    """Action model for security operations tasks."""48 49    task_type: TaskType = Field(..., description="The task being performed")50    action_type: ActionType = Field(..., description="Type of action to take")51 52    redacted_text: Optional[str] = Field(None, description="Redacted text for PII task")53    public_resources: Optional[List[str]] = Field(54        None, description="Resources identified as public"55    )56    fixed_resources: Optional[List[str]] = Field(57        None, description="Resources that were fixed"58    )59    ghost_users: Optional[List[str]] = Field(None, description="Ghost users identified")60    disabled_users: Optional[List[str]] = Field(61        None, description="Users that were disabled"62    )63    classification: Optional[str] = Field(64        None, description="Log classification (e.g., MALWARE, TRUE_POSITIVE, etc.)"65    )66    severity: Optional[str] = Field(67        None, description="Severity level (LOW, MEDIUM, HIGH, CRITICAL)"68    )69    log_alerts: Optional[List[Dict[str, Any]]] = Field(70        None, description="Analyzed log alerts with classifications"71    )72    config_issues: Optional[List[Dict[str, Any]]] = Field(73        None, description="Configuration issues identified"74    )75    hardened_config: Optional[str] = Field(76        None, description="Hardened configuration output"77    )78 79    confidence: Optional[float] = Field(80        None, ge=0.0, le=1.0, description="Confidence score"81    )82    reasoning: Optional[str] = Field(None, description="Reasoning for the action")83 84    model_config = ConfigDict(85        use_enum_values=True,86        extra="allow",  # Allow extra fields for flexibility87    )88 89 90class SecOpsObservation(BaseModel):91    """Observation model for security operations tasks."""92 93    task_type: TaskType = Field(..., description="Current task type")94    task_difficulty: TaskDifficulty = Field(..., description="Task difficulty level")95    objective: str = Field(..., description="Clear objective description")96 97    context: Dict[str, Any] = Field(98        default_factory=dict, description="Scenario context data"99    )100    available_actions: List[str] = Field(101        default_factory=list, description="Actions available"102    )103    current_state: Dict[str, Any] = Field(104        default_factory=dict, description="Current state"105    )106 107    partial_progress: float = Field(108        0.0, ge=0.0, le=1.0, description="Partial completion score"109    )110    step_count: int = Field(0, ge=0, description="Current step number")111    max_steps: int = Field(10, description="Maximum steps allowed")112 113    feedback: Optional[str] = Field(None, description="Feedback on last action")114    detected_issues: List[str] = Field(115        default_factory=list, description="Issues detected so far"116    )117    fixed_issues: List[str] = Field(118        default_factory=list, description="Issues fixed so far"119    )120 121    reward: float = Field(0.0, description="Reward for this step")122    done: bool = Field(False, description="Episode completed flag")123    success: bool = Field(False, description="Task success flag")124    metadata: Dict[str, Any] = Field(125        default_factory=dict, description="Additional metadata"126    )127 128    model_config = ConfigDict(use_enum_values=True, extra="allow")129 130 131class StepResult(BaseModel):132    """Result from a step() call."""133 134    observation: SecOpsObservation135    reward: float = Field(..., description="Reward for this step")136    done: bool = Field(..., description="Episode done flag")137    info: Dict[str, Any] = Field(default_factory=dict, description="Additional info")138 139 140class EpisodeState(BaseModel):141    """Episode state tracking (renamed to avoid conflict with OpenEnv State)."""142 143    episode_id: str144    step_count: int = 0145    task_type: Optional[str] = None146    task_data: Dict[str, Any] = Field(default_factory=dict)147 148    model_config = ConfigDict(use_enum_values=True, extra="allow")149 150 151class TaskConfig(BaseModel):152    """Configuration for a specific task."""153 154    task_type: TaskType155    difficulty: TaskDifficulty156    objective: str157    context: Dict[str, Any]158    success_criteria: Dict[str, Any]159    max_steps: int = 10160 161    model_config = ConfigDict(use_enum_values=True)162