bothari01/secops-env

1"""2SecOps Environment - Core Environment Logic.3 4Security operations environment implementing step/reset/state APIs.5Compatible with OpenEnv framework.6"""7 8import random9from typing import Any, Dict, Optional10from uuid import uuid411 12#0.01 = 1e-913 14 15def _normalize_score(score: float) -> float:16    """Normalize score to be strictly between 0 and 1."""17    if score <= 0:18        return 0.0119    if score >= 1:20        return 0.9921    return score22 23 24from secops_env.models import (25    SecOpsAction,26    SecOpsObservation,27    EpisodeState,28    TaskType,29    TaskDifficulty,30    ActionType,31)32from secops_env.server.tasks.pii_redaction import PIIRedactionTask33from secops_env.server.tasks.public_access import PublicAccessTask34from secops_env.server.tasks.ghost_user import GhostUserTask35from secops_env.server.tasks.log_analysis import LogAnalysisTask36from secops_env.server.tasks.config_hardening import ConfigHardeningTask37from secops_env.server.graders.pii_grader import PIIGrader38from secops_env.server.graders.access_grader import AccessGrader39from secops_env.server.graders.user_grader import UserGrader40from secops_env.server.graders.log_grader import LogGrader41from secops_env.server.graders.config_grader import ConfigGrader42 43 44class SecOpsEnvironment:45    """46    Security Operations Environment for OpenEnv.47 48    Simulates real-world security operations tasks:49    - PII Redaction: Detect and redact personally identifiable information50    - Fix Public Access: Identify and fix overly permissive cloud storage51    - Disable Ghost User: Find and disable orphaned/inactive accounts52 53    Example:54        >>> env = SecOpsEnvironment()55        >>> obs = env.reset(task="pii_redaction")56        >>> print(obs.objective)57        >>> obs = env.step(SecOpsAction(...))58        >>> print(obs.reward)59    """60 61    TASK_REGISTRY = {62        TaskType.PII_REDACTION: PIIRedactionTask,63        TaskType.PUBLIC_ACCESS: PublicAccessTask,64        TaskType.GHOST_USER: GhostUserTask,65        TaskType.LOG_ANALYSIS: LogAnalysisTask,66        TaskType.CONFIG_HARDENING: ConfigHardeningTask,67    }68 69    GRADER_REGISTRY = {70        TaskType.PII_REDACTION: PIIGrader,71        TaskType.PUBLIC_ACCESS: AccessGrader,72        TaskType.GHOST_USER: UserGrader,73        TaskType.LOG_ANALYSIS: LogGrader,74        TaskType.CONFIG_HARDENING: ConfigGrader,75    }76 77    def __init__(self):78        """Initialize the SecOps environment."""79        self._state = EpisodeState(episode_id=str(uuid4()), step_count=0)80        self._current_task: Optional[object] = None81        self._current_task_type: Optional[TaskType] = None82        self._task_data: Dict[str, Any] = {}83        self._reward_history: list[float] = []84        self._reset_count = 085 86    def reset(87        self,88        task: Optional[str] = None,89        difficulty: Optional[str] = None,90        seed: Optional[int] = None,91        **kwargs,92    ) -> SecOpsObservation:93        """94        Reset the environment for a new episode.95 96        Args:97            task: Task type override98            difficulty: Difficulty level override99            seed: Random seed for reproducibility100            **kwargs: Additional options101 102        Returns:103            Initial observation104        """105        if seed is not None:106            random.seed(seed)107 108        self._state = EpisodeState(episode_id=str(uuid4()), step_count=0)109        self._reset_count += 1110        self._reward_history = []111 112        task_type = TaskType(task) if task else random.choice(list(TaskType))113        self._current_task_type = task_type114 115        task_class = self.TASK_REGISTRY[task_type]116        self._current_task = task_class(difficulty=difficulty)117 118        self._task_data = self._current_task.generate_scenario()119        self._state.task_type = task_type.value120        self._state.task_data = self._task_data121 122        observation = self._build_observation(123            reward_accumulated=0.01,124            feedback="Environment ready. Begin security operations.",125        )126 127        return observation128 129    def step(self, action: SecOpsAction) -> SecOpsObservation:130        """131        Execute a step in the environment.132 133        Args:134            action: SecOpsAction to execute135 136        Returns:137            Observation after executing the action138        """139        self._state.step_count += 1140 141        action_task_type = (142            TaskType(action.task_type)143            if isinstance(action.task_type, str)144            else action.task_type145        )146 147        if action_task_type != self._current_task_type:148            return self._build_observation(149                done=False,150                feedback=f"Task mismatch. Current task: {self._current_task_type}, Action task: {action_task_type}",151            )152 153        grader_class = self.GRADER_REGISTRY[self._current_task_type]154        grader = grader_class()155 156        step_reward, feedback, done, success = self._current_task.execute_action(157            action=action, grader=grader, task_data=self._task_data158        )159 160        self._reward_history.append(step_reward)161        # Fix applied here: ensure the accumulated reward is strictly bounded (0, 1)162        reward_accumulated = _normalize_score(sum(self._reward_history))163 164        observation = self._build_observation(165            reward_accumulated=reward_accumulated,166            feedback=feedback,167            done=done,168            success=success,169        )170 171        return observation172 173    def _build_observation(174        self,175        reward_accumulated: float = 0.0,176        feedback: Optional[str] = None,177        done: bool = False,178        success: bool = False,179    ) -> SecOpsObservation:180        """Build observation from current state."""181        task_info = self._current_task.get_info() if self._current_task else {}182 183        partial_progress = self._calculate_partial_progress()184 185        observation = SecOpsObservation(186            task_type=self._current_task_type or TaskType.PII_REDACTION,187            task_difficulty=task_info.get("difficulty", TaskDifficulty.EASY),188            objective=task_info.get("objective", "Complete the security task"),189            context=self._task_data,190            available_actions=[a.value for a in ActionType],191            current_state=self._current_task.get_state() if self._current_task else {},192            partial_progress=partial_progress,193            step_count=self._state.step_count,194            max_steps=self._current_task.max_steps if self._current_task else 10,195            feedback=feedback,196            detected_issues=task_info.get("detected_issues", []),197            fixed_issues=task_info.get("fixed_issues", []),198            reward=reward_accumulated,199            done=done,200            success=success,201            metadata={202                "task_type": str(self._current_task_type)203                if self._current_task_type204                else None,205                "step_count": self._state.step_count,206                "success": success,207            },208        )209 210        return observation211 212    def _calculate_partial_progress(self) -> float:213        """Calculate partial progress toward task completion."""214        if not self._current_task:215            return 0.01216 217        task_info = self._current_task.get_info()218        fixed = len(task_info.get("fixed_issues", []))219        total = (220            self._current_task.total_issues221            if hasattr(self._current_task, "total_issues")222            else 1223        )224 225        return _normalize_score(min(0.99, fixed / max(1, total)))226 227    @property228    def state(self) -> EpisodeState:229        """Get current environment state."""230        return self._state231 232    def get_reward(self) -> float:233        """Get accumulated reward."""234        # Fix applied here: ensure getting the reward also strictly bounds (0, 1)235        return _normalize_score(sum(self._reward_history))