Support updating SSH public key with az aks update --ssh-key-value (#5464)

Author: bingosummer

src/aks-preview/HISTORY.rst

@@ -12,6 +12,11 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+0.5.111
++++++++
+
+* Support updating SSH public key with `az aks update --ssh-key-value`.
+
 0.5.110
 +++++++
 

src/aks-preview/azext_aks_preview/_help.py

@@ -894,6 +894,10 @@
         - name: --cluster-snapshot-id
           type: string
           short-summary: The source cluster snapshot id is used to update existing cluster.
+        - name: --ssh-key-value
+          type: string
+          short-summary: Public key path or key contents to install on node VMs for SSH access. For example,
+                         'ssh-rsa AAAAB...snip...UcyupgH azureuser@linuxvm'.
     examples:
       - name: Reconcile the cluster back to its current state.
         text: az aks update -g MyResourceGroup -n MyManagedCluster

src/aks-preview/azext_aks_preview/_params.py

@@ -116,6 +116,7 @@
     validate_snapshot_name,
     validate_spot_max_price,
     validate_ssh_key,
+    validate_ssh_key_for_update,
     validate_taints,
     validate_user,
     validate_vm_set_type,
@@ -388,6 +389,7 @@ def load_arguments(self, _):
         c.argument('disable_defender', action='store_true', validator=validate_defender_disable_and_enable_parameters)
         c.argument('enable_defender', action='store_true')
         c.argument('defender_config', validator=validate_defender_config_parameter)
+        c.argument('ssh_key_value', type=file_type, completer=FilesCompleter(), validator=validate_ssh_key_for_update)
         # addons
         c.argument('enable_secret_rotation', action='store_true')
         c.argument('disable_secret_rotation', action='store_true')

src/aks-preview/azext_aks_preview/_validators.py

@@ -64,6 +64,20 @@ def validate_ssh_key(namespace):
     namespace.ssh_key_value = content
 
 
+def validate_ssh_key_for_update(namespace):
+    string_or_file = namespace.ssh_key_value
+    if not string_or_file:
+        return
+    content = string_or_file
+    if os.path.exists(string_or_file):
+        logger.info('Use existing SSH public key file: %s', string_or_file)
+        with open(string_or_file, 'r') as f:
+            content = f.read()
+    elif not keys.is_valid_ssh_rsa_public_key(content):
+        raise InvalidArgumentValueError('An RSA key file or key value must be supplied to SSH Key Value')
+    namespace.ssh_key_value = content
+
+
 def validate_create_parameters(namespace):
     if not namespace.name:
         raise CLIError('--name has no value')

src/aks-preview/azext_aks_preview/custom.py

@@ -810,6 +810,7 @@ def aks_update(
     enable_vpa=False,
     disable_vpa=False,
     cluster_snapshot_id=None,
+    ssh_key_value=None,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -14,6 +14,7 @@
 from azure.cli.command_modules.acs._helpers import (
     check_is_msi_cluster,
     format_parameter_name_to_option_name,
+    safe_list_get,
     safe_lower,
 )
 from azure.cli.command_modules.acs._validators import (
@@ -2080,6 +2081,24 @@ def get_disable_vpa(self) -> bool:
         """
         return self._get_disable_vpa(enable_validation=True)
 
+    def get_ssh_key_value_for_update(self) -> Tuple[str, bool]:
+        """Obtain the value of ssh_key_value for "az aks update".
+
+        Note: no_ssh_key will not be decorated into the `mc` object.
+
+        If the user provides a string-like input for --ssh-key-value, the validator function "validate_ssh_key_for_update" will
+        check whether it is a file path, if so, read its content and return; if it is a valid public key, return it.
+        Otherwise, raise error.
+
+        :return: ssh_key_value of string type
+        """
+        # read the original value passed by the command
+        ssh_key_value = self.raw_param.get("ssh_key_value")
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return ssh_key_value
+
 
 class AKSPreviewManagedClusterCreateDecorator(AKSManagedClusterCreateDecorator):
     def __init__(
@@ -2988,6 +3007,27 @@ def update_creation_data(self, mc: ManagedCluster) -> ManagedCluster:
                 source_resource_id=snapshot_id
             )
             mc.creation_data = creation_data
+
+        return mc
+
+    def update_linux_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update Linux profile for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        ssh_key_value = self.context.get_ssh_key_value_for_update()
+
+        if ssh_key_value:
+            mc.linux_profile.ssh = self.models.ContainerServiceSshConfiguration(
+                public_keys=[
+                    self.models.ContainerServiceSshPublicKey(
+                        key_data=ssh_key_value
+                    )
+                ]
+            )
+
         return mc
 
     def update_mc_profile_preview(self) -> ManagedCluster:
@@ -3032,5 +3072,7 @@ def update_mc_profile_preview(self) -> ManagedCluster:
         mc = self.update_vpa(mc)
         # update creation data
         mc = self.update_creation_data(mc)
+        # update linux profile
+        mc = self.update_linux_profile(mc)
 
         return mc