bothari01/secops-env

1"""2SecOps Environment Client.3 4Client for connecting to a SecOps Environment server.5"""6 7import httpx8from typing import Optional, Dict, Any9 10from secops_env.models import SecOpsAction, SecOpsObservation, StepResult11 12 13class SecOpsEnv:14    """15    Client for the SecOps Environment.16 17    Provides access to security operations tasks:18    - pii_redaction: PII detection and redaction19    - public_access: Cloud storage security20    - ghost_user: Account lifecycle management21 22    Example:23        >>> # Sync usage (recommended)24        >>> env = SecOpsEnv(base_url="http://localhost:8000")25        >>> result = env.reset(task="pii_redaction")26        >>> print(result.observation.objective)27        >>> result = env.step(SecOpsAction(...))28        >>> print(result.reward)29        >>> env.close()30    """31 32    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0):33        """Initialize the client with base URL."""34        self.base_url = base_url.rstrip("/")35        self.timeout = timeout36        self._client = httpx.Client(base_url=self.base_url, timeout=timeout)37 38    def reset(39        self,40        task: Optional[str] = None,41        difficulty: Optional[str] = None,42        seed: Optional[int] = None,43        **kwargs,44    ) -> StepResult:45        """46        Reset the environment for a new episode.47 48        Args:49            task: Task type ("pii_redaction", "public_access", "ghost_user")50            difficulty: Difficulty level ("easy", "medium", "hard")51            seed: Random seed for reproducibility52            **kwargs: Additional options53 54        Returns:55            StepResult with initial observation56        """57        params = {}58        if task is not None:59            params["task"] = task60        if difficulty is not None:61            params["difficulty"] = difficulty62        if seed is not None:63            params["seed"] = seed64        params.update(kwargs)65 66        response = self._client.post("/reset", json=params)67        response.raise_for_status()68        data = response.json()69 70        obs_data = data.get("observation", {})71        return StepResult(72            observation=SecOpsObservation(**obs_data),73            reward=data.get("reward", 0.01),74            done=data.get("done", False),75            info=data.get("info", {}),76        )77 78    def step(self, action: SecOpsAction) -> StepResult:79        """80        Execute an action in the environment.81 82        Args:83            action: SecOpsAction to execute84 85        Returns:86            StepResult with observation, reward, and done flag87        """88        payload = action.model_dump()89        response = self._client.post("/step", json={"action": payload})90        response.raise_for_status()91        data = response.json()92 93        obs_data = data.get("observation", {})94        return StepResult(95            observation=SecOpsObservation(**obs_data),96            reward=data.get("reward", 0.01),97            done=data.get("done", False),98            info=data.get("info", {}),99        )100 101    def get_state(self) -> Dict[str, Any]:102        """103        Get current environment state.104 105        Returns:106            Dictionary with episode metadata107        """108        response = self._client.get("/state")109        response.raise_for_status()110        return response.json()111 112    def close(self):113        """Close the HTTP client."""114        self._client.close()115 116    def __enter__(self):117        """Context manager entry."""118        return self119 120    def __exit__(self, exc_type, exc_val, exc_tb):121        """Context manager exit."""122        self.close()123