[Network] az network watcher packet-capture create: Add VMSS suppor…

…t in packet capture (#23649)

src/azure-cli/azure/cli/command_modules/network/_help.py

@@ -6822,6 +6822,14 @@
     short-summary: Name or ID of the VM to target.
   - name: --filters
     short-summary: JSON encoded list of packet filters. Use `@{path}` to load from file.
+  - name: --target
+    short-summary: Name or ID of the Target Resource. If TargetType is AzureVMSS, then target field is mandatory.
+  - name: --target-type
+    short-summary: Resource type of Target. Only AzureVM and AzureVMSS are supported now. Default Value is AzureVM.
+  - name: --include
+    short-summary: Space-separated list of VMSS Instances to include in Packet capture like 0 1 2
+  - name: --exclude
+    short-summary: Space-separated list of VMSS Instances to exclude in Packet capture
 examples:
   - name: Create a packet capture session on a VM.
     text: az network watcher packet-capture create -g MyResourceGroup -n MyPacketCaptureName --vm MyVm --storage-account MyStorageAccount
@@ -6850,6 +6858,18 @@
                 { \\
                     "protocol":"UDP" \\
                 }]'
+  - name: Create a packet capture session on a VMSS.
+    text: |
+        az network watcher packet-capture create -g MyResourceGroup -n MyPacketCaptureName --vm MyVmVMSS \\
+            --storage-account MyStorageAccount --target-type "AzureVMSS"
+  - name: Create a packet capture session on a VMSS with including particular instances.
+    text: |
+        az network watcher packet-capture create -g MyResourceGroup -n MyPacketCaptureName --target MyVmVMSS \\
+            --storage-account MyStorageAccount --target-type "AzureVMSS" --include "0" "1"
+  - name: Create a packet capture session on a VMSS with excluding particular instances.
+    text: |
+        az network watcher packet-capture create -g MyResourceGroup -n MyPacketCaptureName --vm MyVmVMSS \\
+            --storage-account MyStorageAccount --target-type "AzureVMSS" --exclude "0" "1"
 """
 
 helps['network watcher packet-capture delete'] = """

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

@@ -23,7 +23,7 @@
     validate_peering_type, validate_dns_record_type, validate_route_filter, validate_target_listener,
     validate_private_ip_address,
     get_servers_validator, get_public_ip_validator, get_nsg_validator, get_subnet_validator,
-    get_network_watcher_from_vm, get_network_watcher_from_location, validate_capture_size_and_limit,
+    get_network_watcher_from_vm, get_network_watcher_for_pcap_creation, get_network_watcher_from_location, validate_capture_size_and_limit,
     get_asg_validator, get_vnet_validator, validate_ip_tags, validate_ddos_name_or_id,
     validate_service_endpoint_policy, validate_delegations, validate_subresource_list,
     validate_er_peer_circuit, validate_ag_address_pools, validate_custom_error_pages,
@@ -62,7 +62,7 @@ def load_arguments(self, _):
      ExpressRouteLinkMacSecCipher,
      ConnectionMonitorEndpointFilterType, ConnectionMonitorTestConfigurationProtocol,
      PreferredIPVersion, HTTPConfigurationMethod, OutputType, DestinationPortBehavior, CoverageLevel, EndpointType, GatewayLoadBalancerTunnelProtocol,
-     GatewayLoadBalancerTunnelInterfaceType, VpnNatRuleType, VpnNatRuleMode, LoadBalancerBackendAddressAdminState) = self.get_models(
+     GatewayLoadBalancerTunnelInterfaceType, VpnNatRuleType, VpnNatRuleMode, LoadBalancerBackendAddressAdminState, PacketCaptureTargetType) = self.get_models(
          'ApplicationGatewayFirewallMode', 'ApplicationGatewayProtocol', 'ApplicationGatewayRedirectType',
          'ApplicationGatewayRequestRoutingRuleType', 'ApplicationGatewaySkuName', 'ApplicationGatewaySslProtocol', 'AuthenticationMethod',
          'Direction', 'VpnAuthenticationType',
@@ -74,7 +74,7 @@ def load_arguments(self, _):
          'ExpressRouteLinkMacSecCipher',
          'ConnectionMonitorEndpointFilterType', 'ConnectionMonitorTestConfigurationProtocol',
          'PreferredIPVersion', 'HTTPConfigurationMethod', 'OutputType', 'DestinationPortBehavior', 'CoverageLevel', 'EndpointType', 'GatewayLoadBalancerTunnelProtocol',
-         'GatewayLoadBalancerTunnelInterfaceType', 'VpnNatRuleType', 'VpnNatRuleMode', 'LoadBalancerBackendAddressAdminState')
+         'GatewayLoadBalancerTunnelInterfaceType', 'VpnNatRuleType', 'VpnNatRuleMode', 'LoadBalancerBackendAddressAdminState', 'PacketCaptureTargetType')
 
     ZoneType = self.get_models('ZoneType', resource_type=ResourceType.MGMT_NETWORK_DNS)
 
@@ -1674,7 +1674,7 @@ def load_arguments(self, _):
             c.argument('nic', help='Name or ID of the NIC resource to test. If the VM has multiple NICs and IP forwarding is enabled on any of them, this parameter is required.')
 
     with self.argument_context('network watcher packet-capture create') as c:
-        c.argument('watcher_name', ignore_type, validator=get_network_watcher_from_vm)
+        c.argument('watcher_name', ignore_type, validator=get_network_watcher_for_pcap_creation)
         c.ignore('location')
         c.ignore('watcher_rg')
         c.argument('capture_limit', type=int, validator=validate_capture_size_and_limit, help='The maximum size in bytes of the capture output.')
@@ -1683,6 +1683,10 @@ def load_arguments(self, _):
         c.argument('vm', help='Name or ID of the VM to target. If the name of the VM is provided, the --resource-group is required.')
         c.argument('resource_group_name', help='Name of the resource group the target VM is in.')
         c.argument('nic', help='Name or ID of the NIC resource to test. If the VM has multiple NICs and IP forwarding is enabled on any of them, this parameter is required.')
+        c.argument('target_type', help='Target Resource Type, only \'AzureVM\' and \'AzureVMSS\' are supported now', arg_type=get_enum_type(PacketCaptureTargetType))
+        c.argument('target', help='Name or ID of the target, it could be virtual machine or virtual machine scale sets')
+        c.argument('include', nargs='+', help='Space-separated list of VMSS Instances to include in Packet capture like 0 1 2')
+        c.argument('exclude', nargs='+', help='Space-separated list of VMSS Instances to exclude in Packet capture')
 
     with self.argument_context('network watcher test-connectivity') as c:
         c.argument('source_port', type=int)

src/azure-cli/azure/cli/command_modules/network/_validators.py

@@ -1262,6 +1262,13 @@ def load_cert_validator(namespace):
     return load_cert_validator
 
 
+def get_network_watcher_for_pcap_creation(cmd, namespace):
+    if namespace.target_type and namespace.target_type.lower() == "azurevmss":
+        get_network_watcher_from_vmss(cmd, namespace)
+    else:
+        get_network_watcher_from_vm(cmd, namespace)
+
+
 def get_network_watcher_from_vm(cmd, namespace):
     from msrestazure.tools import parse_resource_id
 
@@ -1272,6 +1279,16 @@ def get_network_watcher_from_vm(cmd, namespace):
     get_network_watcher_from_location()(cmd, namespace)
 
 
+def get_network_watcher_from_vmss(cmd, namespace):
+    from msrestazure.tools import parse_resource_id
+
+    compute_client = get_mgmt_service_client(cmd.cli_ctx, ResourceType.MGMT_COMPUTE).virtual_machine_scale_sets
+    vmss_name = parse_resource_id(namespace.target)['name']
+    vmss = compute_client.get(namespace.resource_group_name, vmss_name)
+    namespace.location = vmss.location  # pylint: disable=no-member
+    get_network_watcher_from_location()(cmd, namespace)
+
+
 def get_network_watcher_from_resource(cmd, namespace):
     from azure.cli.core.commands.arm import get_arm_resource_by_id
     resource = get_arm_resource_by_id(cmd.cli_ctx, namespace.resource)
@@ -1738,7 +1755,6 @@ def process_nw_topology_namespace(cmd, namespace):
 
 def process_nw_packet_capture_create_namespace(cmd, namespace):
     from msrestazure.tools import is_valid_resource_id, resource_id
-    get_network_watcher_from_vm(cmd, namespace)
 
     storage_usage = CLIError('usage error: --storage-account NAME_OR_ID [--storage-path '
                              'PATH] [--file-path PATH] | --file-path PATH')
@@ -1748,13 +1764,24 @@ def process_nw_packet_capture_create_namespace(cmd, namespace):
     if namespace.storage_path and not namespace.storage_account:
         raise storage_usage
 
-    if not is_valid_resource_id(namespace.vm):
-        namespace.vm = resource_id(
-            subscription=get_subscription_id(cmd.cli_ctx),
-            resource_group=namespace.resource_group_name,
-            namespace='Microsoft.Compute',
-            type='virtualMachines',
-            name=namespace.vm)
+    if namespace.target_type and namespace.target_type.lower() == "azurevmss":
+        get_network_watcher_from_vmss(cmd, namespace)
+        if not is_valid_resource_id(namespace.target):
+            namespace.target = resource_id(
+                subscription=get_subscription_id(cmd.cli_ctx),
+                resource_group=namespace.resource_group_name,
+                namespace='Microsoft.Compute',
+                type='virtualMachineScaleSets',
+                name=namespace.target)
+    else:
+        get_network_watcher_from_vm(cmd, namespace)
+        if not is_valid_resource_id(namespace.vm):
+            namespace.vm = resource_id(
+                subscription=get_subscription_id(cmd.cli_ctx),
+                resource_group=namespace.resource_group_name,
+                namespace='Microsoft.Compute',
+                type='virtualMachines',
+                name=namespace.vm)
 
     if namespace.storage_account and not is_valid_resource_id(namespace.storage_account):
         namespace.storage_account = resource_id(

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

@@ -6245,18 +6245,26 @@ def show_nw_security_view(cmd, client, resource_group_name, vm, watcher_rg, watc
     return client.begin_get_vm_security_rules(watcher_rg, watcher_name, security_group_view_parameters)
 
 
-def create_nw_packet_capture(cmd, client, resource_group_name, capture_name, vm,
-                             watcher_rg, watcher_name, location=None,
+def create_nw_packet_capture(cmd, client, resource_group_name, capture_name,
+                             watcher_rg, watcher_name, vm=None, location=None,
                              storage_account=None, storage_path=None, file_path=None,
-                             capture_size=None, capture_limit=None, time_limit=None, filters=None):
+                             capture_size=None, capture_limit=None, time_limit=None, filters=None,
+                             target_type=None, target=None, include=None, exclude=None):
     PacketCapture, PacketCaptureStorageLocation = cmd.get_models('PacketCapture', 'PacketCaptureStorageLocation')
+    PacketCaptureMachineScope = cmd.get_models('PacketCaptureMachineScope')
+    # Set the appropriate fields if target is VM
+    pcap_scope = None
+    if not target_type or target_type.lower() != "azurevmss":
+        target = vm
+    else:
+        pcap_scope = PacketCaptureMachineScope(include=include, exclude=exclude)
 
     storage_settings = PacketCaptureStorageLocation(storage_id=storage_account,
                                                     storage_path=storage_path, file_path=file_path)
-    capture_params = PacketCapture(target=vm, storage_location=storage_settings,
+    capture_params = PacketCapture(target=target, storage_location=storage_settings,
                                    bytes_to_capture_per_packet=capture_size,
                                    total_bytes_per_session=capture_limit, time_limit_in_seconds=time_limit,
-                                   filters=filters)
+                                   filters=filters, target_type=target_type, scope=pcap_scope)
     return client.begin_create(watcher_rg, watcher_name, capture_name, capture_params)