[AKS] Support configuring SSH for AKS nodepools

src/azure-cli/azure/cli/command_modules/acs/_consts.py

@@ -67,6 +67,11 @@
 CONST_GPU_DRIVER_INSTALL = "Install"
 CONST_GPU_DRIVER_NONE = "None"
 
+# SSH Access Consts
+CONST_SSH_ACCESS_LOCALUSER = "localuser"
+CONST_SSH_ACCESS_DISABLED = "disabled"
+CONST_SSH_ACCESS_ENTRAID = "entraid"
+
 # consts for ManagedCluster
 # load balancer sku
 CONST_LOAD_BALANCER_SKU_BASIC = "basic"

src/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -61,6 +61,9 @@
     CONST_WEEKINDEX_LAST,
     CONST_MANAGED_CLUSTER_SKU_NAME_BASE,
     CONST_MANAGED_CLUSTER_SKU_NAME_AUTOMATIC,
+    CONST_SSH_ACCESS_LOCALUSER,
+    CONST_SSH_ACCESS_DISABLED,
+    CONST_SSH_ACCESS_ENTRAID,
     CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IP,
     CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IP_CONFIGURATION,
     CONST_AZURE_SERVICE_MESH_INGRESS_MODE_EXTERNAL,
@@ -205,6 +208,11 @@
 node_os_skus_update = [CONST_OS_SKU_AZURELINUX, CONST_OS_SKU_AZURELINUX3, CONST_OS_SKU_AZURECONTAINERLINUX, CONST_OS_SKU_UBUNTU, CONST_OS_SKU_UBUNTU2204, CONST_OS_SKU_UBUNTU2404]
 scale_down_modes = [CONST_SCALE_DOWN_MODE_DELETE, CONST_SCALE_DOWN_MODE_DEALLOCATE]
 pod_ip_allocation_modes = [CONST_NETWORK_POD_IP_ALLOCATION_MODE_DYNAMIC_INDIVIDUAL, CONST_NETWORK_POD_IP_ALLOCATION_MODE_STATIC_BLOCK]
+ssh_accesses = [
+    CONST_SSH_ACCESS_LOCALUSER,
+    CONST_SSH_ACCESS_DISABLED,
+    CONST_SSH_ACCESS_ENTRAID,
+]
 
 # consts for ManagedCluster
 load_balancer_skus = [CONST_LOAD_BALANCER_SKU_BASIC, CONST_LOAD_BALANCER_SKU_STANDARD]
@@ -689,6 +697,13 @@ def load_arguments(self, _):
         c.argument('enable_cost_analysis', action='store_true')
         c.argument('enable_vtpm', action="store_true")
         c.argument('enable_secure_boot', action="store_true")
+        # in creation scenario, use "localuser" as default
+        c.argument(
+            'ssh_access',
+            arg_type=get_enum_type(ssh_accesses),
+            default=CONST_SSH_ACCESS_LOCALUSER,
+            help='Configure SSH access for the first system node pool.',
+        )
         # advanced networking
         c.argument('enable_acns', action='store_true')
         c.argument('disable_acns_observability', action='store_true')
@@ -1291,6 +1306,13 @@ def load_arguments(self, _):
         c.argument("message_of_the_day", validator=validate_message_of_the_day)
         c.argument('enable_vtpm', action='store_true')
         c.argument('enable_secure_boot', action='store_true')
+        # in creation scenario, use "localuser" as default
+        c.argument(
+            'ssh_access',
+            arg_type=get_enum_type(ssh_accesses),
+            default=CONST_SSH_ACCESS_LOCALUSER,
+            help='Configure SSH access for the node pool.',
+        )
         c.argument("if_match")
         c.argument("if_none_match")
         c.argument('gpu_driver', arg_type=get_enum_type(gpu_driver_install_modes))
@@ -1328,6 +1350,13 @@ def load_arguments(self, _):
         c.argument('disable_vtpm', action='store_true')
         c.argument('enable_secure_boot', action='store_true')
         c.argument('disable_secure_boot', action='store_true')
+        # in update scenario, use empty str as default
+        c.argument(
+            'ssh_access',
+            arg_type=get_enum_type(ssh_accesses),
+            help='Update SSH access for the node pool. This change takes effect after the node pool is upgraded.',
+        )
+        c.argument('yes', options_list=['--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')
         c.argument("if_match")
         c.argument("if_none_match")
         c.argument('localdns_config', help='Path to a JSON file to configure the local DNS profile for a new nodepool.')

src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py

@@ -51,6 +51,7 @@
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.util import get_file_json, sdk_no_wait, read_file_content
 from knack.log import get_logger
+from knack.prompting import prompt_y_n
 
 logger = get_logger(__name__)
 
@@ -1990,6 +1991,18 @@ def get_workload_runtime(self) -> Union[str, None]:
         # this parameter does not need validation
         return workload_runtime
 
+    def get_ssh_access(self) -> Union[str, None]:
+        """Obtain the value of ssh_access.
+        """
+        return self.raw_param.get("ssh_access")
+
+    def get_yes(self) -> bool:
+        """Obtain the value of yes.
+
+        :return: bool
+        """
+        return self.raw_param.get("yes")
+
 
 class AKSAgentPoolAddDecorator:
     def __init__(
@@ -2328,6 +2341,20 @@ def set_up_agentpool_security_profile(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def set_up_ssh_access(self, agentpool: AgentPool) -> AgentPool:
+        """Set up the SSH access property for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        ssh_access = self.context.get_ssh_access()
+        if ssh_access is not None:
+            if agentpool.security_profile is None:
+                agentpool.security_profile = self.models.AgentPoolSecurityProfile()  # pylint: disable=no-member
+            agentpool.security_profile.ssh_access = ssh_access
+        return agentpool
+
     def set_up_agentpool_network_profile(self, agentpool: AgentPool) -> AgentPool:
         self._ensure_agentpool(agentpool)
 
@@ -2538,6 +2565,8 @@ def construct_agentpool_profile_default(self, bypass_restore_defaults: bool = Fa
         agentpool = self.set_up_crg_id(agentpool)
         # set up agentpool security profile
         agentpool = self.set_up_agentpool_security_profile(agentpool)
+        # set up agentpool ssh access
+        agentpool = self.set_up_ssh_access(agentpool)
         # set up message of the day
         agentpool = self.set_up_motd(agentpool)
         # set up gpu profile
@@ -2951,6 +2980,31 @@ def update_vtpm(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def update_ssh_access(self, agentpool: AgentPool) -> AgentPool:
+        """Update the SSH access property for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        ssh_access = self.context.get_ssh_access()
+        if ssh_access is not None:
+            if agentpool.security_profile is None:
+                agentpool.security_profile = self.models.AgentPoolSecurityProfile()  # pylint: disable=no-member
+            current_ssh_access = agentpool.security_profile.ssh_access
+            # already set to the same value, directly return
+            if current_ssh_access is not None and current_ssh_access.lower() == ssh_access.lower():
+                return agentpool
+
+            msg = (
+                f"You're going to update agentpool {agentpool.name} ssh access to '{ssh_access}' "
+                "This change will take effect after you upgrade the nodepool. Proceed?"
+            )
+            if not self.context.get_yes() and not prompt_y_n(msg, default="n"):
+                raise DecoratorEarlyExitException()
+            agentpool.security_profile.ssh_access = ssh_access
+        return agentpool
+
     def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool:
         self._ensure_agentpool(agentpool)
 
@@ -3010,6 +3064,8 @@ def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -
         agentpool = self.update_vtpm(agentpool)
         # update secure boot
         agentpool = self.update_secure_boot(agentpool)
+        # update ssh access
+        agentpool = self.update_ssh_access(agentpool)
         # update local DNS profile
         agentpool = self.update_localdns_profile(agentpool)
         # update gpu profile

src/azure-cli/azure/cli/command_modules/acs/custom.py

@@ -71,6 +71,7 @@
     CONST_SCALE_SET_PRIORITY_REGULAR,
     CONST_SECRET_ROTATION_ENABLED,
     CONST_SPOT_EVICTION_POLICY_DELETE,
+    CONST_SSH_ACCESS_LOCALUSER,
     CONST_VIRTUAL_NODE_ADDON_NAME,
     CONST_VIRTUAL_NODE_SUBNET_NAME,
     DecoratorEarlyExitException,
@@ -1075,6 +1076,8 @@ def aks_create(
     enable_gateway_api=False,
     # user-defined scheduler configuration
     enable_upstream_kubescheduler_user_configuration=False,
+    # ssh access
+    ssh_access=CONST_SSH_ACCESS_LOCALUSER,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()
@@ -3599,6 +3602,8 @@ def aks_agentpool_add(
     # local DNS
     localdns_config=None,
     enable_artifact_streaming=False,
+    # ssh access
+    ssh_access=CONST_SSH_ACCESS_LOCALUSER,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()
@@ -3665,6 +3670,9 @@ def aks_agentpool_update(
     gpu_driver=None,
     enable_artifact_streaming=False,
     disable_artifact_streaming=False,
+    # ssh access
+    ssh_access=None,
+    yes=False,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()

src/azure-cli/azure/cli/command_modules/acs/managed_cluster_decorator.py

@@ -6989,6 +6989,11 @@ def get_bootstrap_artifact_source(self) -> Union[str, None]:
         """
         return self.raw_param.get("bootstrap_artifact_source")
 
+    def get_ssh_access(self) -> Union[str, None]:
+        """Obtain the value of ssh_access.
+        """
+        return self.raw_param.get("ssh_access")
+
     def get_bootstrap_container_registry_resource_id(self) -> Union[str, None]:
         """Obtain the value of bootstrap_container_registry_resource_id.
         """
@@ -8848,6 +8853,19 @@ def set_up_node_resource_group_profile(self, mc: ManagedCluster) -> ManagedClust
         mc.node_resource_group_profile = node_resource_group_profile
         return mc
 
+    def set_up_agentpool_profile_ssh_access(self, mc: ManagedCluster) -> ManagedCluster:
+        self._ensure_mc(mc)
+
+        ssh_access = self.context.get_ssh_access()
+        if ssh_access is not None:
+            for agent_pool_profile in (mc.agent_pool_profiles or []):
+                if agent_pool_profile.security_profile is None:
+                    agent_pool_profile.security_profile = (
+                        self.models.AgentPoolSecurityProfile()  # pylint: disable=no-member
+                    )
+                agent_pool_profile.security_profile.ssh_access = ssh_access
+        return mc
+
     def set_up_bootstrap_profile(self, mc: ManagedCluster) -> ManagedCluster:
         self._ensure_mc(mc)
 
@@ -8994,6 +9012,8 @@ def construct_mc_profile_default(self, bypass_restore_defaults: bool = False) ->
         mc = self.set_up_node_resource_group_profile(mc)
         # set up AI toolchain operator
         mc = self.set_up_ai_toolchain_operator(mc)
+        # set up agentpool profile ssh access
+        mc = self.set_up_agentpool_profile_ssh_access(mc)
         # set up bootstrap profile
         mc = self.set_up_bootstrap_profile(mc)
         # set up static egress gateway profile