From 67d455d7946054995f57e81e851a9d19e31fdcc4 Mon Sep 17 00:00:00 2001 From: Nick Carboni Date: Tue, 29 Jul 2025 11:52:27 -0400 Subject: [PATCH 1/3] Add update infraenv function in service client --- service_client/assisted_service_api.py | 24 +++++++++++++ tests/test_assisted_service_api.py | 47 ++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/service_client/assisted_service_api.py b/service_client/assisted_service_api.py index 3757930..0392a6a 100644 --- a/service_client/assisted_service_api.py +++ b/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, diff --git a/tests/test_assisted_service_api.py b/tests/test_assisted_service_api.py index 8b4a79c..7a5966e 100644 --- a/tests/test_assisted_service_api.py +++ b/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.""" From 816b0d444d81116e0414b23e029cf6119bdbe5e7 Mon Sep 17 00:00:00 2001 From: Nick Carboni Date: Tue, 29 Jul 2025 16:22:26 -0400 Subject: [PATCH 2/3] Allow setting ssh public key at cluster creation or via a new tool Resolves https://issues.redhat.com/browse/MGMT-21119 --- server.py | 91 +++++++++++++++++++++++++++--- tests/test_server.py | 129 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 212 insertions(+), 8 deletions(-) diff --git a/server.py b/server.py index f56a77d..9610b61 100644 --- a/server.py +++ b/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.""" diff --git a/tests/test_server.py b/tests/test_server.py index 3a27b1c..72fb7e6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -2,6 +2,8 @@ Unit tests for the server module. """ +# pylint: disable=too-many-lines + import json import os from typing import Generator, Tuple @@ -561,6 +563,55 @@ async def test_create_cluster_success( name, cluster_id="cluster-id", openshift_version=version ) + @pytest.mark.asyncio + async def test_create_cluster_with_ssh_key_success( + self, + mock_inventory_client: Mock, + mock_get_access_token: None, # pylint: disable=unused-argument + ) -> None: + """Test successful create_cluster function with SSH public key.""" + name = "test-cluster" + version = "4.18.2" + base_domain = "example.com" + single_node = False + ssh_public_key = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC test@example.com" + + cluster = create_test_cluster( + cluster_id="cluster-id", + name=name, + openshift_version=version, + ) + infraenv = create_test_infra_env( + infra_env_id="infraenv-id", + name=name, + ) + + mock_inventory_client.create_cluster.return_value = cluster + mock_inventory_client.create_infra_env.return_value = infraenv + + with patch.object( + server, "InventoryClient", return_value=mock_inventory_client + ): + result = await server.create_cluster( + name, version, base_domain, single_node, ssh_public_key + ) + assert result == cluster.id + + mock_inventory_client.create_cluster.assert_called_once_with( + name, + version, + single_node, + base_dns_domain=base_domain, + tags="chatbot", + ssh_public_key=ssh_public_key, + ) + mock_inventory_client.create_infra_env.assert_called_once_with( + name, + cluster_id="cluster-id", + openshift_version=version, + ssh_authorized_key=ssh_public_key, + ) + @pytest.mark.asyncio async def test_set_cluster_vips_success( self, @@ -780,3 +831,81 @@ async def test_cluster_credentials_download_url_zero_expiration( mock_inventory_client.get_presigned_for_cluster_credentials.assert_called_once_with( cluster_id, file_name ) + + @pytest.mark.asyncio + async def test_set_cluster_ssh_key_success( + self, + mock_inventory_client: Mock, + mock_get_access_token: None, # pylint: disable=unused-argument + ) -> None: + """Test successful set_cluster_ssh_key function with cluster and InfraEnvs updated.""" + cluster_id = "test-cluster-id" + ssh_public_key = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC test@example.com" + + # Mock cluster update response + cluster = create_test_cluster(cluster_id=cluster_id) + mock_inventory_client.update_cluster.return_value = cluster + + # Mock InfraEnvs list + mock_infra_envs = [ + {"id": "infraenv-id", "name": "infraenv"}, + ] + mock_inventory_client.list_infra_envs.return_value = mock_infra_envs + + # Mock successful InfraEnv updates + mock_inventory_client.update_infra_env.return_value = None + + with patch.object( + server, "InventoryClient", return_value=mock_inventory_client + ): + result = await server.set_cluster_ssh_key(cluster_id, ssh_public_key) + assert result == cluster.to_str() + + # Verify all expected calls were made + mock_inventory_client.update_cluster.assert_called_once_with( + cluster_id, ssh_public_key=ssh_public_key + ) + mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.update_infra_env.assert_called_once_with( + "infraenv-id", ssh_authorized_key=ssh_public_key + ) + + @pytest.mark.asyncio + async def test_set_cluster_ssh_key_infraenv_failure( + self, + mock_inventory_client: Mock, + mock_get_access_token: None, # pylint: disable=unused-argument + ) -> None: + """Test set_cluster_ssh_key function when InfraEnv update fails.""" + cluster_id = "test-cluster-id" + ssh_public_key = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC test@example.com" + + # Mock cluster update response + cluster = create_test_cluster(cluster_id=cluster_id) + mock_inventory_client.update_cluster.return_value = cluster + + # Mock InfraEnvs list + mock_infra_envs = [ + {"id": "infraenv-id", "name": "infraenv"}, + ] + mock_inventory_client.list_infra_envs.return_value = mock_infra_envs + + # Mock all InfraEnv updates to fail + mock_inventory_client.update_infra_env.side_effect = Exception("Update failed") + + with patch.object( + server, "InventoryClient", return_value=mock_inventory_client + ): + result = await server.set_cluster_ssh_key(cluster_id, ssh_public_key) + # Should return partial failure message + assert "Cluster key updated, but boot image key update failed" in result + assert cluster.to_str() in result + + # Verify all expected calls were made + mock_inventory_client.update_cluster.assert_called_once_with( + cluster_id, ssh_public_key=ssh_public_key + ) + mock_inventory_client.list_infra_envs.assert_called_once_with(cluster_id) + mock_inventory_client.update_infra_env.assert_called_once_with( + "infraenv-id", ssh_authorized_key=ssh_public_key + ) From 438d7b36032ab49c9bf26583ab2b3a522d827156 Mon Sep 17 00:00:00 2001 From: Nick Carboni Date: Tue, 29 Jul 2025 16:27:42 -0400 Subject: [PATCH 3/3] Update doc with ssh key info --- README.md | 9 +++++++++ doc/example_prompts.md | 6 ++++++ 2 files changed, 15 insertions(+) diff --git a/README.md b/README.md index cc505ae..652c4ce 100644 --- a/README.md +++ b/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 diff --git a/doc/example_prompts.md b/doc/example_prompts.md index 4124bb0..33c0a9b 100644 --- a/doc/example_prompts.md +++ b/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`