MGMT-21215: Allow setting or updating cluster ssh public key

README.md

@@ -88,6 +88,7 @@ The MCP server provides the following tools for interacting with the OpenShift A
   * `version`: OpenShift version (string, required)
   * `base_domain`: Base domain for the cluster (string, required)
   * `single_node`: Whether to create a single node cluster (boolean, required)
+  * `ssh_public_key`: SSH public key for accessing cluster nodes (string, optional)
 
 * **install_cluster** - Trigger installation for the assisted installer cluster with the given ID
   * `cluster_id`: Cluster ID (string, required)
@@ -121,6 +122,12 @@ The MCP server provides the following tools for interacting with the OpenShift A
   * `infraenv_id`: Infrastructure environment ID (string, required)
   * `role`: Host role (string, required)
 
+### SSH Key Management
+
+* **set_cluster_ssh_key** - Set or update the SSH public key for a cluster. This allows SSH access to cluster nodes during and after installation.
+  * `cluster_id`: Cluster ID (string, required)
+  * `ssh_public_key`: SSH public key in OpenSSH format (string, required)
+
 ### OpenShift Versions and Operators
 
 * **list_versions** - Lists the available OpenShift versions for installation with the assisted installer
@@ -138,6 +145,8 @@ The MCP server provides the following tools for interacting with the OpenShift A
 * **Create a cluster**: "Create a new cluster named 'my-cluster' with OpenShift 4.14 and base domain 'example.com'"
 * **Check cluster events**: "What events happened on cluster abc123?"
 * **Install a cluster**: "Start the installation for cluster abc123"
+* **Get cluster credentials**: "Get the kubeconfig download link for cluster abc123"
+* **Update SSH key**: "Set the SSH key for cluster abc123 so I can access the nodes"
 
 ## Prometheus Metrics
 

doc/example_prompts.md

@@ -61,6 +61,12 @@ Assign a specific role to a discovered host.
 - "Set host master-node-1 in my-cluster to be a master node"
 - "Assign worker role to host worker-node-2 in my-cluster"
 
+### `set_cluster_ssh_key`
+Set or update the SSH public key for a cluster to allow SSH access to nodes.
+
+**Example Prompts:**
+- "Set the SSH key for my-cluster to ssh-rsa AAAAB3NzaC1yc2E..."
+
 ## Downloads and Resources
 
 ### `cluster_iso_download_url`

server.py

@@ -8,6 +8,7 @@
 import json
 import os
 import asyncio
+from typing import Optional
 
 import requests
 import uvicorn
@@ -289,7 +290,11 @@ async def cluster_iso_download_url(cluster_id: str) -> str:
 @mcp.tool()
 @track_tool_usage()
 async def create_cluster(
-    name: str, version: str, base_domain: str, single_node: bool
+    name: str,
+    version: str,
+    base_domain: str,
+    single_node: bool,
+    ssh_public_key: Optional[str] = None,
 ) -> str:
     """
     Create a new OpenShift cluster.
@@ -306,29 +311,44 @@ async def create_cluster(
         single_node (bool): Whether to create a single-node cluster. Set to True for
             edge deployments or resource-constrained environments. Set to False for
             production high-availability clusters with multiple control plane nodes.
+        ssh_public_key (str, optional): SSH public key for accessing cluster nodes.
+            Providing this key will allow ssh acces to the nodes during and after
+            cluster installation.
 
     Returns:
         str: A the created cluster's id
     """
     log.info(
-        "Creating cluster: name=%s, version=%s, base_domain=%s, single_node=%s",
+        "Creating cluster: name=%s, version=%s, base_domain=%s, single_node=%s, ssh_key_provided=%s",
         name,
         version,
         base_domain,
         single_node,
+        ssh_public_key is not None,
     )
     client = InventoryClient(get_access_token())
-    cluster = await client.create_cluster(
-        name, version, single_node, base_dns_domain=base_domain, tags="chatbot"
-    )
+
+    # Prepare cluster parameters
+    cluster_params = {"base_dns_domain": base_domain, "tags": "chatbot"}
+    if ssh_public_key:
+        cluster_params["ssh_public_key"] = ssh_public_key
+
+    cluster = await client.create_cluster(name, version, single_node, **cluster_params)
     if cluster.id is None:
         log.error("Failed to create cluster %s: cluster ID is unset", name)
         return f"Failed to create cluster {name}: cluster ID is unset"
 
     log.info("Successfully created cluster %s with ID: %s", name, cluster.id)
-    infraenv = await client.create_infra_env(
-        name, cluster_id=cluster.id, openshift_version=cluster.openshift_version
-    )
+
+    # Prepare infra env parameters
+    infraenv_params = {
+        "cluster_id": cluster.id,
+        "openshift_version": cluster.openshift_version,
+    }
+    if ssh_public_key:
+        infraenv_params["ssh_authorized_key"] = ssh_public_key
+
+    infraenv = await client.create_infra_env(name, **infraenv_params)
     log.info(
         "Successfully created InfraEnv for cluster %s with ID: %s",
         cluster.id,
@@ -544,6 +564,61 @@ async def set_host_role(host_id: str, infraenv_id: str, role: str) -> str:
     return result.to_str()
 
 
+@mcp.tool()
+@track_tool_usage()
+async def set_cluster_ssh_key(cluster_id: str, ssh_public_key: str) -> str:
+    """
+    Set or update the SSH public key for a cluster.
+
+    This allows SSH access to cluster nodes during and after installation.
+    Only ISO images downloaded after the update will include the updated key.
+    Discovered hosts should be booted with a new ISO in order to get the new key.
+
+    Args:
+        cluster_id (str): The unique identifier of the cluster to update.
+        ssh_public_key (str): The SSH public key to set for the cluster.
+            This should be a valid SSH public key in OpenSSH format
+            (e.g., 'ssh-rsa AAAAB3NzaC1yc2E... user@host').
+
+    Returns:
+        str: A formatted string containing the updated cluster configuration.
+    """
+    log.info("Setting SSH public key for cluster %s", cluster_id)
+    client = InventoryClient(get_access_token())
+
+    # Update the cluster with the new SSH public key
+    result = await client.update_cluster(cluster_id, ssh_public_key=ssh_public_key)
+    log.info("Successfully updated cluster %s with new SSH key", cluster_id)
+
+    # Get existing InfraEnvs and update them
+    infra_envs = await client.list_infra_envs(cluster_id)
+    log.info("Found %d InfraEnvs for cluster %s", len(infra_envs), cluster_id)
+
+    update_failed = False
+    for infra_env in infra_envs:
+        infra_env_id = infra_env.get("id")
+        if not infra_env_id:
+            log.warning("Skipping InfraEnv without ID: %s", infra_env)
+            continue
+
+        try:
+            await client.update_infra_env(
+                infra_env_id, ssh_authorized_key=ssh_public_key
+            )
+            log.info("Successfully updated InfraEnv %s with new SSH key", infra_env_id)
+        except Exception as e:
+            update_failed = True
+            log.error("Failed to update InfraEnv %s: %s", infra_env_id, str(e))
+
+    if update_failed:
+        return f"Cluster key updated, but boot image key update failed. New cluster: {result.to_str()}"
+
+    log.info(
+        "Successfully updated SSH key for cluster %s and its InfraEnvs", cluster_id
+    )
+    return result.to_str()
+
+
 def list_tools() -> list[str]:
     """List all MCP tools."""
 

service_client/assisted_service_api.py

@@ -282,6 +282,30 @@ async def create_infra_env(
         log.info("Successfully created infrastructure environment '%s'", name)
         return cast(models.InfraEnv, result)
 
+    @sanitize_exceptions
+    async def update_infra_env(
+        self, infra_env_id: str, **update_params: Any
+    ) -> models.InfraEnv:
+        """
+        Update infrastructure environment configuration.
+
+        Args:
+            infra_env_id: The unique identifier of the infrastructure environment to update.
+            **update_params: Infrastructure environment update parameters.
+
+        Returns:
+            models.InfraEnv: The updated infrastructure environment object.
+        """
+        params = models.InfraEnvUpdateParams(**update_params)
+        log.info("Updating infrastructure environment %s", infra_env_id)
+        result = await asyncio.to_thread(
+            self._installer_api().update_infra_env,
+            infra_env_id=infra_env_id,
+            infra_env_update_params=params,
+        )
+        log.info("Successfully updated infrastructure environment %s", infra_env_id)
+        return cast(models.InfraEnv, result)
+
     @sanitize_exceptions
     async def update_cluster(
         self,

tests/test_assisted_service_api.py

@@ -444,6 +444,53 @@ async def test_create_infra_env_success(self, client: InventoryClient) -> None:
             assert infra_env_params.name == name
             assert infra_env_params.pull_secret == client.pull_secret
 
+    @pytest.mark.asyncio
+    async def test_update_infra_env_success(self, client: InventoryClient) -> None:
+        """Test successful update_infra_env function."""
+        infra_env_id = "test-infra-env-id"
+        ssh_key = "ssh-rsa AAAAB3NzaC1yc2E test@example.com"
+
+        infra_env = create_test_infra_env(infra_env_id=infra_env_id)
+
+        with patch.object(
+            client, "_installer_api", return_value=Mock()
+        ) as mock_api_getter:
+            mock_api = mock_api_getter.return_value
+            mock_api.update_infra_env.return_value = infra_env
+
+            result = await client.update_infra_env(
+                infra_env_id, ssh_authorized_key=ssh_key
+            )
+
+            assert result == infra_env
+            mock_api.update_infra_env.assert_called_once()
+
+            # Verify the parameters passed to the API
+            _args, kwargs = mock_api.update_infra_env.call_args
+            assert kwargs["infra_env_id"] == infra_env_id
+            update_params = kwargs["infra_env_update_params"]
+            assert update_params.ssh_authorized_key == ssh_key
+
+    @pytest.mark.asyncio
+    async def test_update_infra_env_api_exception(
+        self, client: InventoryClient
+    ) -> None:
+        """Test update_infra_env function with API exception."""
+        infra_env_id = "test-infra-env-id"
+
+        with patch.object(
+            client, "_installer_api", return_value=Mock()
+        ) as mock_api_getter:
+            mock_api = mock_api_getter.return_value
+            mock_api.update_infra_env.side_effect = ApiException(
+                status=404, reason="Not Found"
+            )
+
+            with pytest.raises(AssistedServiceAPIError, match="API error: Status 404"):
+                await client.update_infra_env(
+                    infra_env_id, ssh_authorized_key="test-key"
+                )
+
     @pytest.mark.asyncio
     async def test_update_cluster_success(self, client: InventoryClient) -> None:
         """Test successful cluster update."""