| import json |
| import threading |
| from typing import Any, Optional |
|
|
| from openenv.core.env_server.interfaces import Environment |
| from openenv.core.env_server.types import State |
|
|
| from cloudenv.models import CloudEnvAction, CloudEnvObservation, CloudEnvState |
| from cloudenv.server.exceptions import CloudEnvBadRequest, CloudEnvNotFound |
| from cloudenv.server.managers.pipeline_manager import PipelineManager |
| from cloudenv.server.managers.reward_manager import RewardManager |
| from cloudenv.server.models.pipeline import Policy |
| from cloudenv.server.providers.aws.adapter import AWSAdapter |
| from cloudenv.server.providers.aws.iam_enforcer import AccessDenied, IAMEnforcer |
| from cloudenv.server.utils.logging import logger |
|
|
|
|
| class CloudEnvironment(Environment): |
| SUPPORTS_CONCURRENT_SESSIONS = False |
|
|
| def __init__( |
| self, |
| aws: AWSAdapter, |
| pipeline_manager: PipelineManager, |
| ) -> None: |
| super().__init__() |
| self._aws_lock = threading.RLock() |
|
|
| self.aws = aws |
| self.pipeline_manager = pipeline_manager |
| self.rewarder = RewardManager(self.pipeline_manager) |
|
|
| self._state: Optional[CloudEnvState] = None |
|
|
| |
| |
| |
| def _sync_iam_to_pipeline(self, action: CloudEnvAction) -> None: |
| pipeline = self.pipeline_manager.get_pipeline(action.task_id) |
| if not pipeline: |
| return |
|
|
| payload = action.payload or {} |
| operation = action.operation |
|
|
| def parse_policy(raw_policy: Any) -> Optional[dict]: |
| if raw_policy is None: |
| return None |
| if isinstance(raw_policy, str): |
| try: |
| return json.loads(raw_policy) |
| except Exception: |
| logger.warning( |
| f"[SYNC] Failed to parse policy: {raw_policy}" |
| ) |
| return None |
| if isinstance(raw_policy, dict): |
| return raw_policy |
| return None |
|
|
| def update_nodes_for_source(role_name: str, policy: dict) -> None: |
| actions = [] |
| for stmt in policy.get("Statement", []): |
| acts = stmt.get("Action", []) |
| if isinstance(acts, str): |
| acts = [acts] |
| actions.extend(acts) |
|
|
| for src_key, nodes in pipeline.broken_graph.items(): |
| _, src_instance = ( |
| src_key.split(":", 1) |
| if ":" in src_key |
| else (src_key, src_key) |
| ) |
|
|
| if src_instance != role_name: |
| continue |
|
|
| for node in nodes: |
| |
| if any(a.startswith(node.service + ":") for a in actions): |
| node.policies = [Policy(**policy)] |
|
|
| def update_nodes_for_resource( |
| service_id: str, instance_id: str, policy: dict |
| ) -> None: |
| for nodes in pipeline.broken_graph.values(): |
| for node in nodes: |
| if ( |
| node.service == service_id |
| and node.instance == instance_id |
| ): |
| node.policies = [Policy(**policy)] |
|
|
| logger.debug(f"[SYNC] Processing {action.service}:{operation}") |
|
|
| if action.service == "iam" and operation == "PutRolePolicy": |
| role_name = payload.get("RoleName") |
| policy_doc = parse_policy(payload.get("PolicyDocument")) |
| if role_name and policy_doc: |
| update_nodes_for_source(role_name, policy_doc) |
|
|
| elif action.service == "s3" and operation == "PutBucketPolicy": |
| bucket = payload.get("Bucket") |
| policy_doc = parse_policy(payload.get("Policy")) |
| if bucket and policy_doc: |
| update_nodes_for_resource("s3", bucket, policy_doc) |
|
|
| elif action.service == "sqs" and operation == "SetQueueAttributes": |
| queue_url = payload.get("QueueUrl") |
| attrs = payload.get("Attributes", {}) or {} |
| policy_doc = parse_policy(attrs.get("Policy")) |
| if queue_url and policy_doc: |
| update_nodes_for_resource("sqs", queue_url, policy_doc) |
|
|
| elif action.service == "sns" and operation == "SetTopicAttributes": |
| topic_arn = payload.get("TopicArn") |
| attr_name = payload.get("AttributeName") |
| policy_doc = parse_policy(payload.get("AttributeValue")) |
| if topic_arn and str(attr_name) == "Policy" and policy_doc: |
| update_nodes_for_resource("sns", topic_arn, policy_doc) |
|
|
| elif action.service == "dynamodb" and operation == "PutResourcePolicy": |
| table = payload.get("ResourceArn") or payload.get("TableName") |
| policy_doc = parse_policy(payload.get("Policy")) |
| if table and policy_doc: |
| update_nodes_for_resource("dynamodb", table, policy_doc) |
|
|
| logger.debug(f"[SYNC] Completed for {action.service}:{operation}") |
|
|
| |
| |
| |
| def reset( |
| self, |
| seed: Optional[int] = None, |
| episode_id: Optional[str] = None, |
| **kwargs: Any, |
| ) -> CloudEnvObservation: |
|
|
| |
| |
|
|
| task_id = kwargs.get("task_id", None) |
| logger.info(f"[RESET] task={task_id}, episode={episode_id}") |
|
|
| self.aws.reset() |
| result = self.aws.get_last_operation_result() |
|
|
| return CloudEnvObservation( |
| episode_id=episode_id, |
| task_id=task_id, |
| step_count=0, |
| operation_result=result, |
| reward=0.0, |
| done=False, |
| current_pipeline_state={}, |
| ) |
|
|
| |
| |
| |
| def step( |
| self, |
| action: CloudEnvAction, |
| timeout_s: Optional[float] = None, |
| **kwargs: Any, |
| ) -> CloudEnvObservation: |
|
|
| logger.info( |
| f"[STEP] task={action.task_id}, episode={action.episode_id}" |
| ) |
|
|
| if action.operation == "InitializeEnvironment": |
| print("REACHED") |
| self._state = CloudEnvState( |
| task_id=action.task_id, |
| episode_id=action.episode_id, |
| step_count=0, |
| reward_history=[], |
| ) |
|
|
| return CloudEnvObservation( |
| episode_id=self._state.episode_id, |
| task_id=self._state.task_id, |
| step_count=self._state.step_count, |
| operation_result={}, |
| reward=0.0, |
| done=False, |
| current_pipeline_state={}, |
| ) |
|
|
| if not action.episode_id or not action.task_id: |
| raise CloudEnvBadRequest("action requires episode_id and task_id") |
|
|
| if action.task_id != self._state.task_id: |
| raise CloudEnvBadRequest("task mismatch") |
|
|
| pipeline = self.pipeline_manager.get_pipeline(action.task_id) |
| if pipeline is None: |
| raise CloudEnvNotFound( |
| f"No pipeline found for task_id={action.task_id}" |
| ) |
|
|
| logger.info( |
| f"[STEP] task={action.task_id}, episode={action.episode_id}, " |
| f"service={action.service}, operation={action.operation}" |
| ) |
|
|
| operation = action.operation |
|
|
| |
| |
| |
| if operation in ["AddConfig", "UpdateConfig", "DeleteConfig"]: |
| logger.debug( |
| f"[CONFIG] operation={operation}, payload={action.payload}" |
| ) |
|
|
| if operation == "AddConfig": |
| pipeline.add_node_config( |
| action.service, |
| action.instance_id, |
| action.payload.get("config", {}), |
| ) |
| elif operation == "UpdateConfig": |
| pipeline.update_node_config( |
| action.service, |
| action.instance_id, |
| action.payload.get("config", {}), |
| ) |
| elif operation == "DeleteConfig": |
| pipeline.delete_node_config( |
| action.service, |
| action.instance_id, |
| action.payload.get("keys", []), |
| ) |
|
|
| reward, done = self.rewarder.compute_reward( |
| pipeline=pipeline, action=action, state=self._state |
| ) |
|
|
| self._state.step_count += 1 |
|
|
| logger.info( |
| f"[STEP RESULT] step={self._state.step_count}, reward={reward}, done={done}" |
| ) |
|
|
| return CloudEnvObservation( |
| episode_id=self._state.episode_id, |
| task_id=self._state.task_id, |
| step_count=self._state.step_count, |
| operation_result={"operation": operation, "status": "SUCCESS"}, |
| reward=reward, |
| done=done, |
| current_pipeline_state=pipeline.broken_graph, |
| ) |
|
|
| |
| |
| |
| with self._aws_lock: |
| success = True |
|
|
| try: |
| logger.debug(f"[AWS EXECUTE] payload={action.payload}") |
|
|
| self.aws.execute( |
| episode_id=action.episode_id, |
| task_id=action.task_id, |
| service=action.service, |
| action=action.operation, |
| payload=action.payload, |
| principal_arn=action.principal_arn, |
| ) |
|
|
| logger.info( |
| f"[AWS SUCCESS] {action.service}:{action.operation}" |
| ) |
|
|
| except AccessDenied as e: |
| success = False |
| logger.warning( |
| f"[ACCESS DENIED] {action.service}:{action.operation} error={str(e)}" |
| ) |
|
|
| self.aws.record_operation_failure( |
| episode_id=action.episode_id, |
| task_id=action.task_id, |
| service=action.service, |
| operation=action.operation, |
| error=str(e), |
| error_type=type(e).__name__, |
| principal=action.principal_arn |
| or IAMEnforcer.DEFAULT_ROOT_PRINCIPAL, |
| ) |
|
|
| except Exception as e: |
| success = False |
| logger.error( |
| f"[AWS ERROR] {action.service}:{action.operation} error={str(e)}" |
| ) |
|
|
| self.aws.record_operation_failure( |
| episode_id=action.episode_id, |
| task_id=action.task_id, |
| service=action.service, |
| operation=action.operation, |
| error=str(e), |
| error_type=type(e).__name__, |
| principal=action.principal_arn |
| or IAMEnforcer.DEFAULT_ROOT_PRINCIPAL, |
| ) |
|
|
| if success: |
| self._sync_iam_to_pipeline(action) |
|
|
| reward, done = self.rewarder.compute_reward( |
| pipeline=pipeline, action=action, state=self._state |
| ) |
|
|
| self._state.step_count += 1 |
| result = self.aws.get_last_operation_result() |
|
|
| logger.info( |
| f"[STEP RESULT] step={self._state.step_count}, reward={reward}, done={done}" |
| ) |
|
|
| return CloudEnvObservation( |
| episode_id=self._state.episode_id, |
| task_id=self._state.task_id, |
| step_count=self._state.step_count, |
| operation_result=result, |
| reward=reward, |
| done=done, |
| current_pipeline_state=pipeline.broken_graph, |
| ) |
|
|
| @property |
| def state(self) -> State: |
| return self._state |
|
|
