| |
|
|
| import uuid |
| from typing import Any, Dict, List, Optional, Set |
|
|
| from cloudenv.server.providers.aws.adapter import AWSAdapter |
| from pydantic import BaseModel, Field |
|
|
|
|
| class PolicyStatement(BaseModel): |
| Effect: str |
| Action: List[str] |
| Resource: str |
|
|
|
|
| class Policy(BaseModel): |
| Version: str |
| Statement: List[PolicyStatement] |
|
|
|
|
| class PipelineNode(BaseModel): |
| |
| service: str |
| instance: str |
| policies: List[Policy] |
| config: Dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class Pipeline: |
| CREATE_ACTIONS = { |
| "s3": "CreateBucket", |
| "lambda": "CreateFunction", |
| "sqs": "CreateQueue", |
| "sns": "CreateTopic", |
| "dynamodb": "CreateTable", |
| "iam": "CreateRole", |
| "ec2": "RunInstances", |
| } |
|
|
| def __init__(self, pipeline_id: Optional[str] = None): |
| self.pipeline_id = pipeline_id or uuid.uuid4().hex |
| self.broken_graph: Dict[str, List[PipelineNode]] = {} |
| self.ideal_graph: Dict[str, List[PipelineNode]] = {} |
|
|
| |
| |
| |
|
|
| def add_from_json( |
| self, json_data: Dict[str, List[Any]], kind: str = "broken" |
| ): |
| graph = self.get_graph(kind) |
|
|
| for src_service, nodes in json_data.items(): |
| if src_service not in graph: |
| graph[src_service] = [] |
|
|
| for node_data in nodes: |
| |
| service_full = getattr(node_data, "service", src_service) |
| instance_name = getattr(node_data, "instance", None) |
|
|
| if ":" in service_full: |
| service_name, parsed_instance = service_full.split(":", 1) |
| instance_name = instance_name or parsed_instance |
| else: |
| service_name = service_full |
| instance_name = instance_name or service_name |
|
|
| |
| |
| |
| policies_raw = getattr(node_data, "policies", []) |
| policies: List[Policy] = [] |
|
|
| for p in policies_raw: |
| if isinstance(p, dict): |
| policies.append(Policy(**p)) |
| else: |
| policies.append(Policy(**p.dict())) |
|
|
| |
| |
| |
| config = getattr(node_data, "config", {}) or {} |
|
|
| |
| |
| |
| existing_node = next( |
| ( |
| n |
| for n in graph[src_service] |
| if n.service == service_name |
| and n.instance == instance_name |
| ), |
| None, |
| ) |
|
|
| if existing_node: |
| |
| for new_policy in policies: |
| if new_policy not in existing_node.policies: |
| existing_node.policies.append(new_policy) |
|
|
| |
| existing_node.config.update(config) |
|
|
| else: |
| |
| node = PipelineNode( |
| service=service_name, |
| instance=instance_name, |
| policies=policies, |
| config=config, |
| ) |
| graph[src_service].append(node) |
|
|
| |
| |
| |
|
|
| def get_graph(self, kind: str) -> Dict[str, List[PipelineNode]]: |
| return self.broken_graph if kind == "broken" else self.ideal_graph |
|
|
| def get_node( |
| self, service: str, instance: str, kind: str = "broken" |
| ) -> Optional[PipelineNode]: |
| graph = self.get_graph(kind) |
| for nodes in graph.values(): |
| for node in nodes: |
| if node.service == service and node.instance == instance: |
| return node |
| return None |
|
|
| def all_nodes(self) -> List[PipelineNode]: |
| nodes: List[PipelineNode] = [] |
| for graph in [self.broken_graph, self.ideal_graph]: |
| for _, dsts in graph.items(): |
| nodes.extend(dsts) |
| return nodes |
|
|
| |
| |
| |
|
|
| def build_payload(self, service: str, instance: str) -> Dict: |
| if service == "s3": |
| return {"Bucket": instance} |
| if service == "sqs": |
| return {"QueueName": instance} |
| if service == "sns": |
| return {"Name": instance} |
| if service == "dynamodb": |
| return { |
| "TableName": instance, |
| "KeySchema": [{"AttributeName": "id", "KeyType": "HASH"}], |
| "AttributeDefinitions": [ |
| {"AttributeName": "id", "AttributeType": "S"} |
| ], |
| "BillingMode": "PAY_PER_REQUEST", |
| } |
| if service == "lambda": |
| return { |
| "FunctionName": instance, |
| "Runtime": "python3.12", |
| "Role": "arn:aws:iam::123456789012:role/dummy", |
| "Handler": "index.handler", |
| "Code": {"ZipFile": b"dummy"}, |
| } |
| if service == "iam": |
| return { |
| "RoleName": instance, |
| "AssumeRolePolicyDocument": "{}", |
| } |
| if service == "ec2": |
| return {"MinCount": 1, "MaxCount": 1} |
| return {} |
|
|
| def materialize(self, aws: AWSAdapter, episode_id: str) -> None: |
| created: Set[str] = set() |
| print("Pipeline materialization started") |
|
|
| for node in self.all_nodes(): |
|
|
| node_key = f"{node.service}:{node.instance}" |
| if node_key in created: |
| continue |
|
|
| action = self.CREATE_ACTIONS.get(node.service) |
| if not action: |
| continue |
|
|
| payload = self.build_payload(node.service, node.instance) |
|
|
| try: |
| aws.execute( |
| episode_id=episode_id, |
| task_id=self.pipeline_id, |
| service=node.service, |
| action=action, |
| payload=payload, |
| ) |
| created.add(node_key) |
|
|
| |
| if node.policies: |
| iam_enforcer = aws.iam_enforcer |
| policies_dicts = [p.model_dump() for p in node.policies] |
| iam_enforcer.set_instance_policies( |
| node.service, node.instance, policies_dicts |
| ) |
|
|
| except Exception as e: |
| print("Materialization error:", e) |
| |
| pass |
|
|
| |
| |
| |
|
|
| def add_node_config( |
| self, |
| service: str, |
| instance: str, |
| config: Dict[str, Any], |
| kind: str = "broken", |
| ) -> bool: |
| node = self.get_node(service, instance, kind) |
| if not node: |
| return False |
| node.config.update(config) |
| return True |
|
|
| def update_node_config( |
| self, |
| service: str, |
| instance: str, |
| config: Dict[str, Any], |
| kind: str = "broken", |
| ) -> bool: |
| return self.add_node_config(service, instance, config, kind) |
|
|
| def delete_node_config( |
| self, |
| service: str, |
| instance: str, |
| keys: List[str], |
| kind: str = "broken", |
| ) -> bool: |
| node = self.get_node(service, instance, kind) |
| if not node: |
| return False |
|
|
| for key in keys: |
| node.config.pop(key, None) |
|
|
| return True |
|
|
| def get_node_config( |
| self, service: str, instance: str, kind: str = "broken" |
| ) -> Optional[Dict[str, Any]]: |
| node = self.get_node(service, instance, kind) |
| return node.config if node else None |
|
|
| def get_all_configs( |
| self, kind: str = "broken" |
| ) -> Dict[str, Dict[str, Any]]: |
| graph = self.get_graph(kind) |
| result = {} |
|
|
| for src_service, nodes in graph.items(): |
| for node in nodes: |
| key = f"{src_service}->{node.service}:{node.instance}" |
| result[key] = node.config |
|
|
| return result |
|
|
