awslabs.eks-mcp-server 0.1.1__py3-none-any.whl

awslabs/eks_mcp_server/k8s_apis.py

@@ -0,0 +1,506 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Kubernetes API client for the EKS MCP Server."""
+
+import base64
+import os
+import tempfile
+from awslabs.eks_mcp_server.models import Operation
+from loguru import logger
+from typing import Any, Dict, List, Optional
+
+
+class K8sApis:
+    """Class for managing Kubernetes API client.
+
+    This class provides a simplified interface for interacting with the Kubernetes API
+    using the official Kubernetes Python client.
+    """
+
+    def __init__(self, endpoint, token, ca_data):
+        """Initialize Kubernetes API client.
+
+        Args:
+            endpoint: Kubernetes API endpoint
+            token: Authentication token
+            ca_data: CA certificate data (base64 encoded) - required for SSL verification
+        """
+        try:
+            from kubernetes import client, dynamic
+
+            configuration = client.Configuration()
+            configuration.host = endpoint
+            configuration.api_key = {'authorization': f'Bearer {token}'}
+
+            # Store the CA cert file path for cleanup
+            self._ca_cert_file_path = None
+
+            # Always enable SSL verification with CA data
+            configuration.verify_ssl = True
+
+            # Create a temporary file for the CA certificate using a context manager
+            try:
+                with tempfile.NamedTemporaryFile(delete=False) as ca_cert_file:
+                    ca_cert_data = base64.b64decode(ca_data)
+                    ca_cert_file.write(ca_cert_data)
+                    # File is automatically closed when exiting the with block
+
+                    # Store the path for cleanup and set the SSL CA cert
+                    self._ca_cert_file_path = ca_cert_file.name
+                    # Set the SSL CA cert to the temporary file path
+                    # Use setattr to avoid potential attribute access issues
+                    setattr(configuration, 'ssl_ca_cert', ca_cert_file.name)
+            except Exception as e:
+                # If we have a path and the file exists, clean it up
+                if (
+                    hasattr(self, '_ca_cert_file_path')
+                    and self._ca_cert_file_path
+                    and os.path.exists(self._ca_cert_file_path)
+                ):
+                    os.unlink(self._ca_cert_file_path)
+                raise e
+
+            # Create base API client
+            self.api_client = client.ApiClient(configuration)
+
+            # Create dynamic client
+            self.dynamic_client = dynamic.DynamicClient(self.api_client)
+
+        except ImportError:
+            logger.error('kubernetes package not installed')
+            raise
+
+    def _patch_resource(
+        self,
+        resource,
+        body: Optional[Dict[str, Any]],
+        name: Optional[str],
+        namespace: Optional[str] = None,
+        **kwargs,
+    ) -> Any:
+        """Patch a resource with strategic merge patch, falling back to merge patch if needed.
+
+        Args:
+            resource: The dynamic resource object
+            body: The resource body to patch with
+            name: Name of the resource
+            namespace: Namespace of the resource (if namespaced)
+            **kwargs: Additional arguments for the API call
+
+        Returns:
+            The API response
+        """
+        try:
+            # First try with strategic merge patch (default)
+            return resource.patch(
+                body=body,
+                name=name,
+                namespace=namespace,
+                content_type='application/strategic-merge-patch+json',
+                **kwargs,
+            )
+        except Exception as e:
+            # If we get a 415 error, try with merge patch
+            if '415' in str(e) or 'Unsupported Media Type' in str(e):
+                logger.warning(
+                    f'Strategic merge patch not supported for {resource.kind}, falling back to merge patch'
+                )
+                return resource.patch(
+                    body=body,
+                    name=name,
+                    namespace=namespace,
+                    content_type='application/merge-patch+json',
+                    **kwargs,
+                )
+            # Re-raise other errors
+            raise
+
+    def manage_resource(
+        self,
+        operation: Operation,
+        kind: str,
+        api_version: str,
+        name: Optional[str] = None,
+        namespace: Optional[str] = None,
+        body: Optional[dict] = None,
+        **kwargs,
+    ) -> Any:
+        """Manage a single Kubernetes resource with the specified operation using dynamic client.
+
+        Args:
+            operation: Operation to perform (Operation.CREATE, Operation.REPLACE, etc.)
+            kind: Resource kind (e.g., 'Pod', 'Service')
+            api_version: API version (e.g., 'v1', 'apps/v1')
+            name: Resource name (required for replace, patch, delete, read)
+            namespace: Namespace of the resource (optional)
+            body: Resource body (required for create, replace, patch)
+            **kwargs: Additional arguments for the API call
+
+        Returns:
+            The API response
+        """
+        # Validate parameters based on operation
+        if (
+            operation in [Operation.REPLACE, Operation.PATCH, Operation.DELETE, Operation.READ]
+            and not name
+        ):
+            raise ValueError(f'Resource name is required for {operation.value} operation')
+
+        if operation in [Operation.CREATE, Operation.REPLACE, Operation.PATCH] and not body:
+            raise ValueError(f'Resource body is required for {operation.value} operation')
+
+        try:
+            # Get the API resource
+            resource = self.dynamic_client.resources.get(api_version=api_version, kind=kind)
+
+            # Set kind and apiVersion in the body if provided
+            if body:
+                body['kind'] = kind
+                body['apiVersion'] = api_version
+
+                # Set name and namespace in metadata if provided
+                if name:
+                    if 'metadata' not in body:
+                        body['metadata'] = {}
+                    body['metadata']['name'] = name
+                if namespace:
+                    if 'metadata' not in body:
+                        body['metadata'] = {}
+                    body['metadata']['namespace'] = namespace
+
+            # Perform the operation based on the operation type
+            if operation == Operation.CREATE:
+                return resource.create(body=body, namespace=namespace, **kwargs)
+            elif operation == Operation.REPLACE:
+                return resource.replace(body=body, name=name, namespace=namespace, **kwargs)
+            elif operation == Operation.PATCH:
+                return self._patch_resource(resource, body, name, namespace, **kwargs)
+            elif operation == Operation.DELETE:
+                return resource.delete(name=name, namespace=namespace, **kwargs)
+            elif operation == Operation.READ:
+                return resource.get(name=name, namespace=namespace, **kwargs)
+            else:
+                raise ValueError(f'Unsupported operation: {operation.value}')
+
+        except Exception as e:
+            # Re-raise with more context
+            raise ValueError(f'Error managing {kind} resource: {str(e)}')
+
+    def list_resources(
+        self,
+        kind: str,
+        api_version: str,
+        namespace: Optional[str] = None,
+        label_selector: Optional[str] = None,
+        field_selector: Optional[str] = None,
+        **kwargs,
+    ) -> Any:
+        """List Kubernetes resources of a specific kind using dynamic client.
+
+        Args:
+            kind: Resource kind (e.g., 'Pod', 'Service')
+            api_version: API version (e.g., 'v1', 'apps/v1')
+            namespace: Namespace to list resources from (optional)
+            label_selector: Label selector to filter resources (optional)
+            field_selector: Field selector to filter resources (optional)
+            **kwargs: Additional arguments for the API call
+
+        Returns:
+            The API response containing the list of resources
+        """
+        try:
+            # Get the API resource
+            resource = self.dynamic_client.resources.get(api_version=api_version, kind=kind)
+
+            # Prepare kwargs for the list operation
+            list_kwargs = {}
+            if label_selector:
+                list_kwargs['label_selector'] = label_selector
+            if field_selector:
+                list_kwargs['field_selector'] = field_selector
+
+            # Add any additional kwargs
+            list_kwargs.update(kwargs)
+
+            # List resources
+            if namespace:
+                return resource.get(namespace=namespace, **list_kwargs)
+            else:
+                return resource.get(**list_kwargs)
+
+        except Exception as e:
+            # Re-raise with more context
+            raise ValueError(f'Error listing {kind} resources: {str(e)}')
+
+    def apply_from_yaml(
+        self, yaml_objects: list, namespace: str = 'default', force: bool = True, **kwargs
+    ) -> tuple:
+        """Apply YAML objects to the cluster with support for custom resources and updates.
+
+        This method improves upon the standard create_from_yaml by:
+        1. Supporting custom resources through the dynamic client
+        2. Supporting updates to existing resources when force=True
+
+        Args:
+            yaml_objects: List of YAML objects to apply
+            namespace: Default namespace to use for namespaced resources
+            force: Whether to update resources if they already exist (like kubectl apply)
+            **kwargs: Additional arguments for the API calls
+
+        Returns:
+            Tuple of (results, created_count, updated_count)
+        """
+        results = []
+        created_count = 0
+        updated_count = 0
+
+        for obj in yaml_objects:
+            if not obj:
+                continue
+
+            # Extract key information from the object
+            kind = obj.get('kind')
+            api_version = obj.get('apiVersion')
+            metadata = obj.get('metadata', {})
+            name = metadata.get('name')
+            obj_namespace = metadata.get('namespace', namespace)
+
+            if not kind or not api_version or not name:
+                raise ValueError('Invalid resource: missing kind, apiVersion, or name')
+
+            try:
+                # Get the API resource
+                resource = self.dynamic_client.resources.get(api_version=api_version, kind=kind)
+
+                # Check if resource exists
+                exists = False
+                if force:
+                    try:
+                        resource.get(
+                            name=name, namespace=obj_namespace if resource.namespaced else None
+                        )
+                        exists = True
+                    except Exception:
+                        # Resource doesn't exist, will be created
+                        exists = False
+
+                # Apply the resource
+                if exists and force:
+                    # Update existing resource - use patch only
+                    result = self._patch_resource(
+                        resource,
+                        obj,
+                        name,
+                        obj_namespace if resource.namespaced else None,
+                        **kwargs,
+                    )
+                    updated_count += 1
+                else:
+                    # Create new resource
+                    result = resource.create(
+                        body=obj,
+                        namespace=obj_namespace if resource.namespaced else None,
+                        **kwargs,
+                    )
+                    created_count += 1
+
+                results.append(result)
+
+            except Exception as e:
+                # Add context to the error
+                resource_name = f'{obj_namespace}/{name}' if obj_namespace else name
+                raise ValueError(f'Error applying {kind} {resource_name}: {str(e)}')
+
+        return results, created_count, updated_count
+
+    def get_events(
+        self,
+        kind: str,
+        name: str,
+        namespace: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """Get events related to a specific Kubernetes resource.
+
+        Args:
+            kind: Kind of the involved object (e.g., 'Pod', 'Deployment')
+            name: Name of the involved object
+            namespace: Namespace of the involved object (optional for non-namespaced resources)
+
+        Returns:
+            List of events related to the specified object
+        """
+        try:
+            # Get the Event resource using the dynamic client
+            event_resource = self.dynamic_client.resources.get(api_version='v1', kind='Event')
+
+            # Prepare field selector to filter events
+            field_selector = f'involvedObject.kind={kind},involvedObject.name={name}'
+
+            # If namespace is provided, get events from that namespace
+            # Otherwise, search across all namespaces
+            if namespace:
+                events_response = event_resource.get(
+                    namespace=namespace, field_selector=field_selector
+                )
+            else:
+                events_response = event_resource.get(field_selector=field_selector)
+
+            # Process events
+            result = []
+            for event in events_response.items:
+                # Dynamic client resources always have to_dict()
+                event_dict = event.to_dict()
+
+                # Extract relevant fields and handle camelCase field names
+                first_timestamp = event_dict.get('firstTimestamp')
+                last_timestamp = event_dict.get('lastTimestamp')
+                source = event_dict.get('source', {})
+
+                result.append(
+                    {
+                        'first_timestamp': str(first_timestamp) if first_timestamp else None,
+                        'last_timestamp': str(last_timestamp) if last_timestamp else None,
+                        'count': event_dict.get('count'),
+                        'message': event_dict.get('message', ''),
+                        'reason': event_dict.get('reason'),
+                        'reporting_component': source.get('component'),
+                        'type': event_dict.get('type'),
+                    }
+                )
+
+            return result
+
+        except Exception as e:
+            # Re-raise with more context
+            resource_name = f'{namespace + "/" if namespace else ""}{name}'
+            raise ValueError(f'Error getting events for {kind} {resource_name}: {str(e)}')
+
+    def get_pod_logs(
+        self,
+        pod_name: str,
+        namespace: str,
+        container_name: Optional[str] = None,
+        since_seconds: Optional[int] = None,
+        tail_lines: Optional[int] = None,
+        limit_bytes: Optional[int] = None,
+    ) -> str:
+        """Get logs from a pod.
+
+        Args:
+            pod_name: Name of the pod
+            namespace: Namespace of the pod
+            container_name: Container name (optional, if pod contains more than one container)
+            since_seconds: Only return logs newer than this many seconds (optional)
+            tail_lines: Number of lines to return from the end of the logs (optional)
+            limit_bytes: Maximum number of bytes to return (optional)
+
+        Returns:
+            Pod logs as a string
+        """
+        try:
+            # Get the Pod resource using the dynamic client
+            pod_resource = self.dynamic_client.resources.get(api_version='v1', kind='Pod')
+
+            # Prepare parameters for the log subresource
+            params = {}
+            if container_name:
+                params['container'] = container_name
+            if since_seconds:
+                params['sinceSeconds'] = since_seconds
+            if tail_lines:
+                params['tailLines'] = tail_lines
+            if limit_bytes:
+                params['limitBytes'] = limit_bytes
+
+            # Call the log subresource (note: singular 'log', not 'logs')
+            logs_response = pod_resource.log.get(name=pod_name, namespace=namespace, **params)
+
+            return logs_response
+
+        except Exception as e:
+            # Re-raise with more context
+            raise ValueError(f'Error getting logs from pod {namespace}/{pod_name}: {str(e)}')
+
+    def get_api_versions(self) -> List[str]:
+        """Get preferred API versions from the Kubernetes cluster.
+
+        Returns only the preferred (stable) API version for each group, avoiding alpha/beta versions
+        when stable versions are available.
+
+        Returns:
+            List of preferred API versions (e.g., ['v1', 'apps/v1', 'networking.k8s.io/v1'])
+        """
+        try:
+            from kubernetes import client
+
+            api_versions: set[str] = set()
+
+            # Get core API version (v1)
+            try:
+                core_api = client.CoreApi(self.api_client)
+                core_version_obj = core_api.get_api_versions()
+
+                # Extract versions safely
+                if core_version_obj is not None:
+                    # Try to get versions as a list of strings
+                    versions = getattr(core_version_obj, 'versions', None)
+                    if versions is not None and isinstance(versions, list):
+                        for version in versions:
+                            if isinstance(version, str):
+                                api_versions.add(version)
+            except Exception as e:
+                logger.warning(f'Error getting core API versions: {str(e)}')
+                raise ValueError(f'Error getting API versions: {str(e)}')
+
+            # Get API groups and their preferred versions
+            try:
+                apis_api = client.ApisApi(self.api_client)
+                api_groups_obj = apis_api.get_api_versions()
+
+                # Extract groups safely
+                if api_groups_obj is not None:
+                    groups = getattr(api_groups_obj, 'groups', None)
+                    if groups is not None and isinstance(groups, list):
+                        for group in groups:
+                            if group is not None:
+                                # Try to get preferred version
+                                preferred_version = getattr(group, 'preferred_version', None)
+                                if preferred_version is not None:
+                                    group_version = getattr(
+                                        preferred_version, 'group_version', None
+                                    )
+                                    if group_version is not None and isinstance(
+                                        group_version, str
+                                    ):
+                                        api_versions.add(group_version)
+            except Exception as e:
+                logger.warning(f'Error getting API groups: {str(e)}')
+
+            # Convert to sorted list
+            return sorted(api_versions)
+
+        except Exception as e:
+            # Re-raise with more context
+            raise ValueError(f'Error getting API versions: {str(e)}')
+
+    def __del__(self):
+        """Clean up temporary files when the object is garbage collected."""
+        if (
+            hasattr(self, '_ca_cert_file_path')
+            and self._ca_cert_file_path
+            and os.path.exists(self._ca_cert_file_path)
+        ):
+            try:
+                os.unlink(self._ca_cert_file_path)
+            except Exception:
+                # Ignore errors during cleanup
+                pass

awslabs/eks_mcp_server/k8s_client_cache.py

@@ -0,0 +1,164 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Kubernetes client cache for the EKS MCP Server."""
+
+import base64
+from awslabs.eks_mcp_server.aws_helper import AwsHelper
+from awslabs.eks_mcp_server.k8s_apis import K8sApis
+from cachetools import TTLCache
+
+
+# Presigned url timeout in seconds
+URL_TIMEOUT = 60
+TOKEN_PREFIX = 'k8s-aws-v1.'
+K8S_AWS_ID_HEADER = 'x-k8s-aws-id'
+
+# 14 minutes in seconds (buffer before the 15-minute token expiration)
+TOKEN_TTL = 14 * 60
+
+
+class K8sClientCache:
+    """Singleton class for managing Kubernetes API client cache.
+
+    This class provides a centralized cache for Kubernetes API clients
+    to avoid creating multiple clients for the same cluster.
+    """
+
+    # Singleton instance
+    _instance = None
+
+    def __new__(cls):
+        """Ensure only one instance of K8sClientCache exists."""
+        if cls._instance is None:
+            cls._instance = super(K8sClientCache, cls).__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        """Initialize the K8s client cache."""
+        # Only initialize once
+        if hasattr(self, '_initialized') and self._initialized:
+            return
+
+        # Client cache with TTL to handle token expiration
+        self._client_cache = TTLCache(maxsize=100, ttl=TOKEN_TTL)
+
+        # Clients for credential retrieval
+        self._eks_client = None
+        self._sts_client = None
+
+        self._initialized = True
+
+    def _get_eks_client(self):
+        """Get or create the EKS client."""
+        if self._eks_client is None:
+            self._eks_client = AwsHelper.create_boto3_client('eks')
+        return self._eks_client
+
+    def _get_sts_client(self):
+        """Get or create the STS client with event handlers registered."""
+        if self._sts_client is None:
+            sts_client = AwsHelper.create_boto3_client('sts')
+
+            # Register STS event handlers
+            sts_client.meta.events.register(
+                'provide-client-params.sts.GetCallerIdentity',
+                self._retrieve_k8s_aws_id,
+            )
+            sts_client.meta.events.register(
+                'before-sign.sts.GetCallerIdentity',
+                self._inject_k8s_aws_id_header,
+            )
+
+            self._sts_client = sts_client
+
+        return self._sts_client
+
+    def _retrieve_k8s_aws_id(self, params, context, **kwargs):
+        """Retrieve the Kubernetes AWS ID from parameters."""
+        if K8S_AWS_ID_HEADER in params:
+            context[K8S_AWS_ID_HEADER] = params.pop(K8S_AWS_ID_HEADER)
+
+    def _inject_k8s_aws_id_header(self, request, **kwargs):
+        """Inject the Kubernetes AWS ID header into the request."""
+        if K8S_AWS_ID_HEADER in request.context:
+            request.headers[K8S_AWS_ID_HEADER] = request.context[K8S_AWS_ID_HEADER]
+
+    def _get_cluster_credentials(self, cluster_name: str):
+        """Get credentials for an EKS cluster (private method).
+
+        Args:
+            cluster_name: Name of the EKS cluster
+
+        Returns:
+            Tuple of (endpoint, token, ca_data)
+
+        Raises:
+            ValueError: If the cluster credentials are invalid
+            Exception: If there's an error getting the cluster credentials
+        """
+        eks_client = self._get_eks_client()
+        sts_client = self._get_sts_client()
+
+        # Get cluster details
+        response = eks_client.describe_cluster(name=cluster_name)
+        endpoint = response['cluster']['endpoint']
+        ca_data = response['cluster']['certificateAuthority']['data']
+
+        # Generate a presigned URL for authentication
+        url = sts_client.generate_presigned_url(
+            'get_caller_identity',
+            Params={K8S_AWS_ID_HEADER: cluster_name},
+            ExpiresIn=URL_TIMEOUT,
+            HttpMethod='GET',
+        )
+
+        # Create the token from the presigned URL
+        token = TOKEN_PREFIX + base64.urlsafe_b64encode(url.encode('utf-8')).decode(
+            'utf-8'
+        ).rstrip('=')
+
+        return endpoint, token, ca_data
+
+    def get_client(self, cluster_name: str) -> K8sApis:
+        """Get a Kubernetes client for the specified cluster.
+
+        This is the only public method to access K8s API clients.
+
+        Args:
+            cluster_name: Name of the EKS cluster
+
+        Returns:
+            K8sApis instance
+
+        Raises:
+            ValueError: If the cluster credentials are invalid
+            Exception: If there's an error getting the cluster credentials
+        """
+        if cluster_name not in self._client_cache:
+            try:
+                # Create a new client
+                endpoint, token, ca_data = self._get_cluster_credentials(cluster_name)
+
+                # Validate credentials
+                if not endpoint or not token or endpoint is None or token is None:
+                    raise ValueError('Invalid cluster credentials')
+
+                self._client_cache[cluster_name] = K8sApis(endpoint, token, ca_data)
+            except ValueError:
+                # Re-raise ValueError for invalid credentials
+                raise
+            except Exception as e:
+                # Re-raise any other exceptions
+                raise Exception(f'Failed to get cluster credentials: {str(e)}')
+
+        return self._client_cache[cluster_name]