from .base_agent import BaseAgent from ..aws_connector import AWSConnector import json class IAMAgent(BaseAgent): """ An agent specialized in handling AWS IAM tasks. """ def execute(self, command: str, **kwargs): """ Executes a given command related to IAM. Includes a confirmation step for destructive actions. """ if command == 'list_users': return self._list_users() elif command == 'create_user': return self._create_user(**kwargs) elif command == 'delete_user': user_name = kwargs.get('user_name', '') print(f"WARNING: You are about to delete the IAM user '{user_name}'. This is irreversible.") confirm = input("Are you sure you want to proceed? (yes/no): ") if confirm.lower() == 'yes': return self._delete_user(**kwargs) else: return {"status": "cancelled", "message": "Delete user command cancelled by user."} elif command == 'create_eks_cluster_role': role_name = kwargs.get('role_name', '') print(f"You are about to create an IAM role '{role_name}' for an EKS cluster.") confirm = input("Are you sure you want to proceed? (yes/no): ") if confirm.lower() == 'yes': return self._create_eks_cluster_role(**kwargs) else: return {"status": "cancelled", "message": "Create EKS cluster role command cancelled by user."} elif command == 'create_eks_nodegroup_role': role_name = kwargs.get('role_name', '') print(f"You are about to create an IAM role '{role_name}' for an EKS node group.") confirm = input("Are you sure you want to proceed? (yes/no): ") if confirm.lower() == 'yes': return self._create_eks_nodegroup_role(**kwargs) else: return {"status": "cancelled", "message": "Create EKS nodegroup role command cancelled by user."} elif command == 'create_lambda_execution_role': role_name = kwargs.get('role_name', '') print(f"You are about to create an IAM role '{role_name}' for a Lambda function.") confirm = input("Are you sure you want to proceed? (yes/no): ") if confirm.lower() == 'yes': return self._create_lambda_execution_role(**kwargs) else: return {"status": "cancelled", "message": "Create Lambda execution role command cancelled by user."} elif command == 'list_roles': return self._list_roles() elif command == 'list_policies': return self._list_policies() else: raise NotImplementedError(f"Command '{command}' is not supported by IAMAgent.") def _list_users(self): """ Lists all IAM users. """ print("IAMAgent: Listing IAM users...") try: iam_client = AWSConnector.get_client('iam') response = iam_client.list_users() users = [{'UserName': user['UserName'], 'Arn': user['Arn']} for user in response['Users']] return {"status": "success", "message": f"Found {len(users)} IAM users.", "users": users} except Exception as e: print(f"Error listing IAM users: {e}") return {"status": "error", "message": str(e)} def _create_user(self, user_name: str): """ Creates a new IAM user. """ # ... (code is unchanged) def _delete_user(self, user_name: str): """ Deletes an IAM user. """ # ... (code is unchanged) def _create_eks_cluster_role(self, role_name: str): """ Creates an IAM role for an EKS cluster and attaches the necessary policies. """ print(f"IAMAgent: Creating EKS cluster role '{role_name}'...") try: iam_client = AWSConnector.get_client('iam') # Define the trust policy for EKS trust_policy = { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": "eks.amazonaws.com" }, "Action": "sts:AssumeRole" } ] } # Create the role create_role_response = iam_client.create_role( RoleName=role_name, AssumeRolePolicyDocument=json.dumps(trust_policy), Description='IAM role for EKS cluster management.' ) role_arn = create_role_response['Role']['Arn'] # Attach the AmazonEKSClusterPolicy iam_client.attach_role_policy( RoleName=role_name, PolicyArn='arn:aws:iam::aws:policy/AmazonEKSClusterPolicy' ) # Attach the AmazonEC2ContainerRegistryReadOnly policy iam_client.attach_role_policy( RoleName=role_name, PolicyArn='arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly' ) print(f"EKS cluster role '{role_name}' created and policies attached.") return {"status": "success", "message": f"EKS cluster role '{role_name}' created and policies attached.", "role_arn": role_arn} except iam_client.exceptions.EntityAlreadyExistsException: print(f"IAM role '{role_name}' already exists. Fetching ARN...") role_arn = iam_client.get_role(RoleName=role_name)['Role']['Arn'] return {"status": "success", "message": f"EKS cluster role '{role_name}' already exists.", "role_arn": role_arn} except Exception as e: print(f"Error creating EKS cluster role: {e}") return {"status": "error", "message": str(e)} def _create_eks_nodegroup_role(self, role_name: str): """ Creates an IAM role for an EKS node group and attaches the necessary policies. """ # ... (code is unchanged) def _create_lambda_execution_role(self, role_name: str): """ Creates an IAM role for a Lambda function and attaches the AWSLambdaBasicExecutionRole policy. """ print(f"IAMAgent: Creating Lambda execution role '{role_name}'...") try: iam_client = AWSConnector.get_client('iam') # Define the trust policy for Lambda trust_policy = { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": "lambda.amazonaws.com" }, "Action": "sts:AssumeRole" } ] } # Create the role create_role_response = iam_client.create_role( RoleName=role_name, AssumeRolePolicyDocument=json.dumps(trust_policy), Description='IAM role for Lambda function execution.' ) role_arn = create_role_response['Role']['Arn'] # Attach the AWSLambdaBasicExecutionRole policy iam_client.attach_role_policy( RoleName=role_name, PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' ) return {"status": "success", "message": f"Lambda execution role '{role_name}' created and policy attached.", "role_arn": role_arn} except Exception as e: print(f"Error creating Lambda execution role: {e}") return {"status": "error", "message": str(e)} def _list_roles(self): """ Lists all IAM roles. """ print("IAMAgent: Listing IAM roles...") try: iam_client = AWSConnector.get_client('iam') response = iam_client.list_roles() roles = [{'RoleName': role['RoleName'], 'Arn': role['Arn']} for role in response['Roles']] return {"status": "success", "message": f"Found {len(roles)} IAM roles.", "roles": roles} except Exception as e: print(f"Error listing IAM roles: {e}") return {"status": "error", "message": str(e)} def _list_policies(self): """ Lists all IAM policies (customer managed). """ print("IAMAgent: Listing IAM policies...") try: iam_client = AWSConnector.get_client('iam') response = iam_client.list_policies(Scope='Local') # Local for customer managed policies policies = [{'PolicyName': policy['PolicyName'], 'Arn': policy['Arn']} for policy in response['Policies']] return {"status": "success", "message": f"Found {len(policies)} IAM policies.", "policies": policies} except Exception as e: print(f"Error listing IAM policies: {e}") return {"status": "error", "message": str(e)}