Add option for aks node restriction

src/aks-preview/HISTORY.rst

@@ -11,6 +11,7 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+*  Add --enable-node-restriction/--disable-node-restriction to enable/disable node restriction feature
 
 0.5.94
 ++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -805,6 +805,12 @@
         - name: --defender-config
           type: string
           short-summary: Path to JSON file containing Microsoft Defender profile configurations.
+        - name: --enable-node-restriction
+          type: bool
+          short-summary: Enable node restriction option on cluster.
+        - name: --disable-node-restriction
+          type: bool
+          short-summary: Disable node restriction option on cluster.
     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

@@ -320,6 +320,7 @@ def load_arguments(self, _):
         c.argument('apiserver_subnet_id', validator=validate_apiserver_subnet_id, is_preview=True)
         c.argument('dns_zone_resource_id')
         c.argument('enable_keda', action='store_true', is_preview=True)
+        c.argument('enable_node_restriction', action='store_true', is_preview=True, help="enable node restriction for cluster")
         # nodepool
         c.argument('host_group_id', validator=validate_host_group_id, is_preview=True)
         c.argument('crg_id', validator=validate_crg_id, is_preview=True)
@@ -414,6 +415,8 @@ def load_arguments(self, _):
         c.argument('apiserver_subnet_id', validator=validate_apiserver_subnet_id, is_preview=True)
         c.argument('enable_keda', action='store_true', is_preview=True)
         c.argument('disable_keda', action='store_true', is_preview=True)
+        c.argument('enable_node_restriction', action='store_true', is_preview=True, help="enable node restriction for cluster")
+        c.argument('disable_node_restriction', action='store_true', is_preview=True, help="disable node restriction for cluster")
 
     with self.argument_context('aks upgrade') as c:
         c.argument('kubernetes_version', completer=get_k8s_upgrades_completion_list)

src/aks-preview/azext_aks_preview/custom.py

@@ -667,6 +667,7 @@ def aks_create(
     apiserver_subnet_id=None,
     dns_zone_resource_id=None,
     enable_keda=False,
+    enable_node_restriction=False,
     # nodepool
     host_group_id=None,
     crg_id=None,
@@ -783,6 +784,8 @@ def aks_update(
     apiserver_subnet_id=None,
     enable_keda=False,
     disable_keda=False,
+    enable_node_restriction=False,
+    disable_node_restriction=False,
 ):
     # 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

@@ -76,6 +76,7 @@
 ManagedClusterStorageProfileSnapshotController = TypeVar('ManagedClusterStorageProfileSnapshotController')
 ManagedClusterIngressProfileWebAppRouting = TypeVar("ManagedClusterIngressProfileWebAppRouting")
 ManagedClusterSecurityProfileDefender = TypeVar("ManagedClusterSecurityProfileDefender")
+ManagedClusterSecurityProfileNodeRestriction = TypeVar("ManagedClusterSecurityProfileNodeRestriction")
 
 
 # pylint: disable=too-few-public-methods
@@ -1395,6 +1396,67 @@ def get_defender_config(self) -> Union[ManagedClusterSecurityProfileDefender, No
         )
         return azure_defender
 
+    def _get_enable_node_restriction(self, enable_validation: bool = False) -> bool:
+        """Internal function to obtain the value of enable_node_restriction.
+
+        This function supports the option of enable_node_restriction. When enabled, if both enable_node_restriction and disable_node_restriction are
+        specified, raise a MutuallyExclusiveArgumentError.
+
+        :return: bool
+        """
+        # Read the original value passed by the command.
+        enable_node_restriction = self.raw_param.get("enable_node_restriction")
+
+        # This parameter does not need dynamic completion.
+        if enable_validation:
+            if enable_node_restriction and self._get_disable_node_restriction(enable_validation=False):
+                raise MutuallyExclusiveArgumentError(
+                    "Cannot specify --enable-node-restriction and --disable-node-restriction at the same time."
+                )
+
+        return enable_node_restriction
+
+    def get_enable_node_restriction(self) -> bool:
+        """Obtain the value of enable_node_restriction.
+
+        This function will verify the parameter by default. If both enable_node_restriction and disable_node_restriction are specified, raise a
+        MutuallyExclusiveArgumentError.
+
+        :return: bool
+        """
+        return self._get_enable_node_restriction(enable_validation=True)
+
+    def _get_disable_node_restriction(self, enable_validation: bool = False) -> bool:
+        """Internal function to obtain the value of disable_node_restriction.
+
+        This function supports the option of enable_validation. When enabled, if both enable_node_restriction and disable_node_restriction are
+        specified, raise a MutuallyExclusiveArgumentError.
+
+        :return: bool
+        """
+        # Read the original value passed by the command.
+        disable_node_restriction = self.raw_param.get("disable_node_restriction")
+
+        # This option is not supported in create mode, hence we do not read the property value from the `mc` object.
+        # This parameter does not need dynamic completion.
+        if enable_validation:
+            if disable_node_restriction and self._get_enable_node_restriction(enable_validation=False):
+                raise MutuallyExclusiveArgumentError(
+                    "Cannot specify --enable-node-restriction and --disable-node-restriction at the same time."
+                )
+
+        return disable_node_restriction
+
+    def get_disable_node_restriction(self) -> bool:
+        """Obtain the value of disable_node_restriction.
+
+        This function will verify the parameter by default. If both enable_node_restriction and disable_node_restriction are specified, raise a
+        MutuallyExclusiveArgumentError.
+
+        :return: bool
+        """
+        return self._get_disable_node_restriction(enable_validation=True)
+
 
 class AKSPreviewManagedClusterCreateDecorator(AKSManagedClusterCreateDecorator):
     def __init__(
@@ -1701,6 +1763,22 @@ def set_up_defender(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def set_up_node_restriction(self, mc: ManagedCluster) -> ManagedCluster:
+        """Set up security profile nodeRestriction for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        if self.context.get_enable_node_restriction():
+            if mc.security_profile is None:
+                    mc.security_profile = self.models.ManagedClusterSecurityProfile()
+            mc.security_profile.node_restriction = self.models.ManagedClusterSecurityProfileNodeRestriction(
+                enabled=True,
+            )
+
+        return mc
+
     def construct_mc_profile_preview(self, bypass_restore_defaults: bool = False) -> ManagedCluster:
         """The overall controller used to construct the default ManagedCluster profile.
 
@@ -1729,6 +1807,8 @@ def construct_mc_profile_preview(self, bypass_restore_defaults: bool = False) ->
 
         # set up azure keyvalut kms
         mc = self.set_up_azure_keyvault_kms(mc)
+        # set up node restriction
+        mc = self.set_up_node_restriction(mc)
         # set up cluster snapshot
         mc = self.set_up_creationdata_of_cluster_snapshot(mc)
         # set up storage profile
@@ -2035,6 +2115,33 @@ def update_defender(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def update_node_restriction(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update security profile nodeRestriction for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        if self.context.get_enable_node_restriction():
+            if mc.security_profile is None:
+                mc.security_profile = self.models.ManagedClusterSecurityProfile()
+            if mc.security_profile.node_restriction is None:
+                mc.security_profile.node_restriction = self.models.ManagedClusterSecurityProfileNodeRestriction()
+
+            # set enabled
+            mc.security_profile.node_restriction.enabled = True
+
+        if self.context.get_disable_node_restriction():
+            if mc.security_profile is None:
+                mc.security_profile = self.models.ManagedClusterSecurityProfile()
+            if mc.security_profile.node_restriction is None:
+                mc.security_profile.node_restriction = self.models.ManagedClusterSecurityProfileNodeRestriction()
+
+            # set disabled
+            mc.security_profile.node_restriction.enabled = False
+
+        return mc
+
     def update_mc_profile_preview(self) -> ManagedCluster:
         """The overall controller used to update the preview ManagedCluster profile.
 
@@ -2063,6 +2170,8 @@ def update_mc_profile_preview(self) -> ManagedCluster:
 
         # update azure keyvalut kms
         mc = self.update_azure_keyvault_kms(mc)
+        # update node restriction
+        mc = self.update_node_restriction(mc)
         # update stroage profile
         mc = self.update_storage_profile(mc)
         # update workload auto scaler profile

src/aks-preview/azext_aks_preview/tests/latest/test_aks_commands.py

@@ -260,6 +260,39 @@ def test_aks_create_and_update_with_managed_aad_enable_azure_rbac(self, resource
             self.check('aadProfile.enableAzureRbac', False)
         ])
 
+    @AllowLargeResponse()
+    @AKSCustomResourceGroupPreparer(random_name_length=17, name_prefix='clitest', location='westus2')
+    def test_aks_create_and_update_with_node_restriction(self, resource_group, resource_group_location):
+        aks_name = self.create_random_name('cliakstest', 16)
+        self.kwargs.update({
+            'resource_group': resource_group,
+            'name': aks_name,
+            'ssh_key_value': self.generate_ssh_keys()
+        })
+
+        create_cmd = 'aks create --resource-group={resource_group} --name={name} ' \
+                     '--vm-set-type VirtualMachineScaleSets -c 1 ' \
+                     '--enable-node-restriction ' \
+                     '--ssh-key-value={ssh_key_value} -o json'
+        self.cmd(create_cmd, checks=[
+            self.check('provisioningState', 'Succeeded'),
+            self.check('securityProfile.nodeRestriction.enabled', True),
+        ])
+
+        update_cmd = 'aks update --resource-group={resource_group} --name={name} ' \
+                     '--disable-node-restriction -o json'
+        self.cmd(update_cmd, checks=[
+            self.check('provisioningState', 'Succeeded'),
+            self.check('securityProfile.nodeRestriction.enabled', False)
+        ])
+
+        update_cmd = 'aks update --resource-group={resource_group} --name={name} ' \
+                     '--enable-node-restriction -o json'
+        self.cmd(update_cmd, checks=[
+            self.check('provisioningState', 'Succeeded'),
+            self.check('securityProfile.nodeRestriction.enabled', True)
+        ])
+
     @AllowLargeResponse()
     @AKSCustomResourceGroupPreparer(random_name_length=17, name_prefix='clitest', location='westus2')
     def test_aks_create_with_ingress_appgw_addon(self, resource_group, resource_group_location):