KSingh08/soc2-auditor

1"""SOC 2 Evidence Auditor — Environment Client."""2 3from typing import Any, Dict4 5from openenv.core import EnvClient6from openenv.core.client_types import StepResult7from openenv.core.env_server.types import State8 9from .models import SOC2Action, SOC2Observation10 11 12class SOC2Env(EnvClient[SOC2Action, SOC2Observation, State]):13    """14    Client for the SOC 2 Evidence Auditor environment.15 16    Maintains a persistent WebSocket connection to the environment server,17    enabling efficient multi-step audit interactions.18 19    Example (Docker):20        >>> env = await SOC2Env.from_docker_image("soc2-auditor:latest")21        >>> result = await env.reset(task_id="pr_approval_check")22        >>> result = await env.step(SOC2Action(23        ...     type="INSPECT_FILE", file_name="pull_request_log.json"24        ... ))25        >>> result = await env.step(SOC2Action(26        ...     type="SUBMIT_DECISION", decision="REJECT", reason="MISSING_APPROVAL"27        ... ))28        >>> await env.close()29 30    Example (SEARCH_LOGS for large log files):31        >>> env = await SOC2Env.from_docker_image("soc2-auditor:latest")32        >>> result = await env.reset(task_id="cloudtrail_privileged_access_audit")33        >>> result = await env.step(SOC2Action(34        ...     type="INSPECT_FILE", file_name="hr_terminations.json"35        ... ))36        >>> result = await env.step(SOC2Action(37        ...     type="SEARCH_LOGS",38        ...     file_name="aws_cloudtrail_full_log.json",39        ...     query_field="username",40        ...     query_value="alice_dev",41        ... ))42        >>> result = await env.step(SOC2Action(43        ...     type="SUBMIT_DECISION", decision="REJECT", reason="INCOMPLETE_REVOCATION"44        ... ))45        >>> await env.close()46 47    Example (running server):48        >>> async with SOC2Env(base_url="http://localhost:8000") as env:49        ...     result = await env.reset(task_id="access_revocation_sla")50        ...     obs = result.observation51        ...     print(obs.control_requirement)52        ...     print(obs.available_files)53    """54 55    def _step_payload(self, action: SOC2Action) -> Dict[str, Any]:56        payload: Dict[str, Any] = {"type": action.type}57        if action.file_name is not None:58            payload["file_name"] = action.file_name59        if action.decision is not None:60            payload["decision"] = action.decision61        if action.reason is not None:62            payload["reason"] = action.reason63        if action.query_field is not None:64            payload["query_field"] = action.query_field65        if action.query_value is not None:66            payload["query_value"] = action.query_value67        return payload68 69    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SOC2Observation]:70        obs_data = payload.get("observation", {})71        observation = SOC2Observation(72            task_id=obs_data.get("task_id", ""),73            control_requirement=obs_data.get("control_requirement", ""),74            available_files=obs_data.get("available_files", []),75            inspected_files=obs_data.get("inspected_files", {}),76            audit_status=obs_data.get("audit_status", "IN_PROGRESS"),77            step_reward=obs_data.get("step_reward", 0.0),78            cumulative_reward=obs_data.get("cumulative_reward", 0.0),79            done=payload.get("done", False),80            reward=payload.get("reward"),81            message=obs_data.get("message", ""),82            metadata=obs_data.get("metadata", {}),83        )84        return StepResult(85            observation=observation,86            reward=payload.get("reward"),87            done=payload.get("done", False),88        )89 90    def _parse_state(self, payload: Dict[str, Any]) -> State:91        return State(92            episode_id=payload.get("episode_id"),93            step_count=payload.get("step_count", 0),94        )95