[AKS] Add NSG Control flags for az aks create / `az aks nodepool ad…

…d` / `az aks nodepool update` (#5467)

* Add NSG Control flags

* Fix

* Add aks command test

* Fix aks command test

* Fix ut

* Fix UT

* Fix live test

* Fix typo

src/aks-preview/HISTORY.rst

@@ -12,6 +12,14 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+0.5.110
++++++++
+
+* Add `--nodepool-asg-ids` and `--nodepool-allowed-host-ports` flags for enabling NSGControl. Related commands:
+  * `az aks create`
+  * `az aks nodepool add`
+  * `az aks nodepool update`
+
 0.5.109
 +++++++
 

src/aks-preview/azext_aks_preview/_help.py

@@ -485,6 +485,12 @@
         - name: --enable-vpa
           type: bool
           short-summary: Enable vertical pod autoscaler for cluster.
+        - name: --nodepool-allowed-host-ports
+          type: string
+          short-summary: Expose host ports on the node pool. When specified, format should be a comma-separated list of ranges with protocol, eg. 80/TCP,443/TCP,4000-5000/TCP.
+        - name: --nodepool-asg-ids
+          type: string
+          short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a comma-separated list of IDs.
     examples:
         - name: Create a Kubernetes cluster with an existing SSH public key.
           text: az aks create -g MyResourceGroup -n MyManagedCluster --ssh-key-value /path/to/publickey
@@ -1278,6 +1284,12 @@
         - name: --disable-windows-outbound-nat
           type: bool
           short-summary: Disable Windows OutboundNAT on Windows agent node pool.
+        - name: --allowed-host-ports
+          type: string
+          short-summary: Expose host ports on the node pool. When specified, format should be a comma-separated list of ranges with protocol, eg. 80/TCP,443/TCP,4000-5000/TCP.
+        - name: --asg-ids
+          type: string
+          short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a comma-separated list of IDs.
     examples:
         - name: Create a nodepool in an existing AKS cluster with ephemeral os enabled.
           text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --node-osdisk-type Ephemeral --node-osdisk-size 48
@@ -1369,6 +1381,12 @@
         - name: --aks-custom-headers
           type: string
           short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2
+        - name: --allowed-host-ports
+          type: string
+          short-summary: Expose host ports on the node pool. When specified, format should be a comma-separated list of ranges with protocol, eg. 80/TCP,443/TCP,4000-5000/TCP.
+        - name: --asg-ids
+          type: string
+          short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a comma-separated list of IDs.
     examples:
       - name: Reconcile the nodepool back to its current state.
         text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster

src/aks-preview/azext_aks_preview/_params.py

@@ -128,6 +128,8 @@
     validate_ksm_labels,
     validate_ksm_annotations,
     validate_disable_windows_outbound_nat,
+    validate_allowed_host_ports,
+    validate_application_security_groups,
 )
 
 # candidates for enumeration
@@ -344,6 +346,8 @@ def load_arguments(self, _):
         # no validation for aks create because it already only supports Linux.
         c.argument('enable_custom_ca_trust', action='store_true')
         c.argument('enable_vpa', action='store_true', is_preview=True, help="enable vertical pod autoscaler for cluster")
+        c.argument('nodepool_allowed_host_ports', validator=validate_allowed_host_ports, is_preview=True, help="allowed host ports for agentpool")
+        c.argument('nodepool_asg_ids', validator=validate_application_security_groups, is_preview=True, help="application security groups for agentpool")
 
     with self.argument_context('aks update') as c:
         # managed cluster paramerters
@@ -507,6 +511,8 @@ def load_arguments(self, _):
         c.argument('gpu_instance_profile', arg_type=get_enum_type(gpu_instance_profiles))
         c.argument('enable_custom_ca_trust', action='store_true', validator=validate_enable_custom_ca_trust)
         c.argument('disable_windows_outbound_nat', action='store_true', validator=validate_disable_windows_outbound_nat)
+        c.argument('allowed_host_ports', validator=validate_allowed_host_ports, is_preview=True)
+        c.argument('asg_ids', validator=validate_application_security_groups, is_preview=True)
 
     with self.argument_context('aks nodepool update') as c:
         c.argument('enable_cluster_autoscaler', options_list=[
@@ -526,6 +532,8 @@ def load_arguments(self, _):
         # extensions
         c.argument('enable_custom_ca_trust', action='store_true', validator=validate_enable_custom_ca_trust)
         c.argument('disable_custom_ca_trust', options_list=['--disable-custom-ca-trust', '--dcat'], action='store_true')
+        c.argument('allowed_host_ports', validator=validate_allowed_host_ports, is_preview=True)
+        c.argument('asg_ids', validator=validate_application_security_groups, is_preview=True)
 
     with self.argument_context('aks nodepool upgrade') as c:
         c.argument('max_surge', validator=validate_max_surge)

src/aks-preview/azext_aks_preview/_validators.py

@@ -316,7 +316,8 @@ def validate_load_balancer_idle_timeout(namespace):
 def validate_load_balancer_backend_pool_type(namespace):
     """validate load balancer backend pool type"""
     if namespace.load_balancer_backend_pool_type is not None:
-        if namespace.load_balancer_backend_pool_type not in [CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IP, CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IPCONFIGURATION]:
+        if namespace.load_balancer_backend_pool_type not in [CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IP,
+                                                             CONST_LOAD_BALANCER_BACKEND_POOL_TYPE_NODE_IPCONFIGURATION]:
             raise InvalidArgumentValueError(
                 f"Invalid Load Balancer Backend Pool Type {namespace.load_balancer_backend_pool_type}, supported values are nodeIP and nodeIPConfiguration")
 
@@ -739,3 +740,35 @@ def validate_ksm_annotations(namespace):
     if namespace.ksm_metric_annotations_allow_list is None:
         return
     validate_ksm_parameter(namespace.ksm_metric_annotations_allow_list)
+
+
+def validate_allowed_host_ports(namespace):
+    if hasattr(namespace, "nodepool_allowed_host_ports"):
+        host_ports = namespace.nodepool_allowed_host_ports
+    else:
+        host_ports = namespace.allowed_host_ports
+    if not host_ports:
+        return
+
+    regex = re.compile(r'^((\d+)|(\d+-\d+))/(tcp|udp)$')
+    for port_range in host_ports.split(","):
+        found = regex.findall(port_range)
+        if found:
+            continue
+        raise InvalidArgumentValueError(
+            "--allowed-host-ports must be a comma-separated list of port ranges in the format of <port-range>/<protocol>"
+        )
+
+
+def validate_application_security_groups(namespace):
+    if hasattr((namespace), "nodepool_asg_ids"):
+        asg_ids = namespace.nodepool_asg_ids
+    else:
+        asg_ids = namespace.asg_ids
+    if not asg_ids:
+        return
+
+    from msrestazure.tools import is_valid_resource_id
+    for asg in asg_ids.split(","):
+        if not is_valid_resource_id(asg):
+            raise InvalidArgumentValueError(asg + " is not a valid Azure resource ID.")

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -34,6 +34,7 @@
 # type variables
 AgentPool = TypeVar("AgentPool")
 AgentPoolsOperations = TypeVar("AgentPoolsOperations")
+PortRange = TypeVar("PortRange")
 
 
 # pylint: disable=too-few-public-methods
@@ -260,6 +261,49 @@ def get_disable_windows_outbound_nat(self) -> bool:
         """
         return self._get_disable_windows_outbound_nat()
 
+    def get_asg_ids(self) -> Union[List[str], None]:
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
+            asg_ids = self.raw_param.get('nodepool_asg_ids')
+        else:
+            asg_ids = self.raw_param.get('asg_ids')
+
+        if asg_ids is None:
+            return None
+        if asg_ids == '':
+            return []
+
+        return asg_ids.split(',')
+
+    def get_allowed_host_ports(self) -> Union[List[PortRange], None]:
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
+            ports = self.raw_param.get('nodepool_allowed_host_ports')
+        else:
+            ports = self.raw_param.get('allowed_host_ports')
+
+        if ports is None:
+            return None
+        if ports == '':
+            return []
+
+        ports = ports.split(',')
+        port_ranges = []
+        import re
+        regex = re.compile(r'^((\d+)|((\d+)-(\d+)))/(tcp|udp)$')
+        for port in ports:
+            r = regex.findall(port)
+            if r[0][1] != '':
+                # single port
+                port_start, port_end = int(r[0][1]), int(r[0][1])
+            else:
+                # port range
+                port_start, port_end = int(r[0][3]), int(r[0][4])
+            port_ranges.append(self.models.PortRange(
+                port_start=port_start,
+                port_end=port_end,
+                protocol=r[0][5].upper(),
+            ))
+        return port_ranges
+
 
 class AKSPreviewAgentPoolAddDecorator(AKSAgentPoolAddDecorator):
     def __init__(
@@ -353,6 +397,18 @@ def set_up_agentpool_windows_profile(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def set_up_agentpool_network_profile(self, agentpool: AgentPool) -> AgentPool:
+        self._ensure_agentpool(agentpool)
+
+        asg_ids = self.context.get_asg_ids()
+        allowed_host_ports = self.context.get_allowed_host_ports()
+        if asg_ids and allowed_host_ports:
+            agentpool.network_profile = self.models.AgentPoolNetworkProfile(
+                application_security_groups=asg_ids,
+                allowed_host_ports=allowed_host_ports,
+            )
+        return agentpool
+
     def construct_agentpool_profile_preview(self) -> AgentPool:
         """The overall controller used to construct the preview AgentPool profile.
 
@@ -374,6 +430,8 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self.set_up_custom_ca_trust(agentpool)
         # set up agentpool windows profile
         agentpool = self.set_up_agentpool_windows_profile(agentpool)
+        # set up agentpool network profile
+        agentpool = self.set_up_agentpool_network_profile(agentpool)
 
         # DO NOT MOVE: keep this at the bottom, restore defaults
         agentpool = self._restore_defaults_in_agentpool(agentpool)
@@ -427,6 +485,19 @@ def update_custom_ca_trust(self, agentpool: AgentPool) -> AgentPool:
             agentpool.enable_custom_ca_trust = False
         return agentpool
 
+    def update_network_profile(self, agentpool: AgentPool) -> AgentPool:
+        self._ensure_agentpool(agentpool)
+
+        asg_ids = self.context.get_asg_ids()
+        allowed_host_ports = self.context.get_allowed_host_ports()
+        if asg_ids or allowed_host_ports:
+            agentpool.network_profile = self.models.AgentPoolNetworkProfile()
+        if asg_ids is not None:
+            agentpool.network_profile.application_security_groups = asg_ids
+        if allowed_host_ports is not None:
+            agentpool.network_profile.allowed_host_ports = allowed_host_ports
+        return agentpool
+
     def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -> AgentPool:
         """The overall controller used to update the preview AgentPool profile.
 
@@ -440,4 +511,8 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
 
         # update custom ca trust
         agentpool = self.update_custom_ca_trust(agentpool)
+
+        # update network profile
+        agentpool = self.update_network_profile(agentpool)
+
         return agentpool

src/aks-preview/azext_aks_preview/custom.py

@@ -627,6 +627,8 @@ def aks_create(
     node_count=3,
     nodepool_tags=None,
     nodepool_labels=None,
+    nodepool_allowed_host_ports=None,
+    nodepool_asg_ids=None,
     node_osdisk_type=None,
     node_osdisk_size=0,
     vm_set_type=None,
@@ -1172,6 +1174,8 @@ def aks_agentpool_add(
     gpu_instance_profile=None,
     enable_custom_ca_trust=False,
     disable_windows_outbound_nat=False,
+    allowed_host_ports=None,
+    asg_ids=None,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()
@@ -1219,6 +1223,8 @@ def aks_agentpool_update(
     # extensions
     enable_custom_ca_trust=False,
     disable_custom_ca_trust=False,
+    allowed_host_ports=None,
+    asg_ids=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

@@ -1948,9 +1948,9 @@ def _get_enable_azure_monitor_metrics(self, enable_validation: bool = False) ->
                     "Cannot specify --enable-azuremonitormetrics and --enable-azuremonitormetrics at the same time."
                 )
             if not check_is_msi_cluster(self.mc):
-                    raise RequiredArgumentMissingError(
-                        "--enable-azuremonitormetrics can only be specified for clusters with managed identity enabled"
-                    )
+                raise RequiredArgumentMissingError(
+                    "--enable-azuremonitormetrics can only be specified for clusters with managed identity enabled"
+                )
         return enable_azure_monitor_metrics
 
     def get_enable_azure_monitor_metrics(self) -> bool: