[Network] az network watcher flow-log: support --vnet --subnet --nic arguments

.vscode/launch.json

@@ -54,4 +54,4 @@
             ]
         }
     ]
-}
+}

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

@@ -6934,6 +6934,12 @@
 parameters:
   - name: --nsg
     short-summary: Name or ID of the Network Security Group to target.
+  - name: --vnet
+    short-summary: Name of the Virtual Network to target.
+  - name: --subnet
+    short-summary: Name or ID of the subnet.
+  - name: --nic
+    short-summary: Name or ID of the Network Interface.
   - name: --enabled
     short-summary: Enable logging.
   - name: --retention
@@ -6949,7 +6955,7 @@
 
 helps['network watcher flow-log create'] = """
 type: command
-short-summary: Create a flow log on a network security group.
+short-summary: Create a flow log on a network security group or virtual network.
 examples:
   - name: Create a flow log with Network Security Group name
     text: >

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

@@ -1694,6 +1694,10 @@ def load_arguments(self, _):
         c.argument('retention', type=int, help='Number of days to retain logs')
         c.argument('storage_account', help='Name or ID of the storage account in which to save the flow logs. '
                                            'Must be in the same region of flow log.')
+
+        c.argument('vnet',options_list=['--vnet'],help='Name or ID of the Virtual Network Resource.')
+        c.argument('subnet',options_list=['--subnet'], help='Name or ID of Subnet')
+        c.argument('nic', options_list=['--nic'],help='Name or ID of the Network Interface (NIC) Resource.')
 
     # temporary solution for compatible with old show command's parameter
     # after old show command's parameter is deprecated and removed,

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

@@ -1498,6 +1498,15 @@ def process_nw_flow_log_create_namespace(cmd, namespace):
         if namespace.traffic_analytics_workspace and not is_valid_resource_id(namespace.traffic_analytics_workspace):
             err_body = '--workspace ID / --workspace NAME --resource-group WORKSPACE_RESOURCE_GROUP'
 
+        if namespace.vnet and not is_valid_resource_id(namespace.vnet):
+            err_body = '--vnet ID / --vnet NAME --resource-group VNET_RESOURCE_GROUP'
+
+        if namespace.subnet and not is_valid_resource_id(namespace.subnet):
+            err_body = '--subnet ID / --subnet NAME --resource-group SUBNET_RESOURCE_GROUP'
+
+        if namespace.nic and not is_valid_resource_id(namespace.nic):
+            err_body = '--nic ID / --nic NAME --resource-group NIC_RESOURCE_GROUP'
+
         if err_body is not None:
             raise CLIError(err_tpl.format(err_body))
 
@@ -1533,6 +1542,33 @@ def process_nw_flow_log_create_namespace(cmd, namespace):
             'name': namespace.traffic_analytics_workspace
         }
         namespace.traffic_analytics_workspace = resource_id(**kwargs)
+
+    if namespace.vnet and not is_valid_resource_id(namespace.vnet):
+        kwargs = {
+            'subscription': get_subscription_id(cmd.cli_ctx),
+            'resource_group': namespace.resource_group_name,
+            'namespace': 'Microsoft.Network',
+            'type': 'virtualNetworks',
+            'name': namespace.vnet
+        }
+
+    if namespace.subnet and not is_valid_resource_id(namespace.subnet):
+        kwargs = {
+            'subscription': get_subscription_id(cmd.cli_ctx),
+            'resource_group': namespace.resource_group_name,
+            'namespace': 'Microsoft.Network',
+            'type': 'virtualNetworks',
+            'name': namespace.subnet
+        }
+
+    if namespace.nic and not is_valid_resource_id(namespace.nic):
+        kwargs = {
+            'subscription': get_subscription_id(cmd.cli_ctx),
+            'resource_group': namespace.resource_group_name,
+            'namespace': 'Microsoft.Network',
+            'type': 'networkInterfaces',
+            'name': namespace.nic
+        }
 
     get_network_watcher_from_location(remove=False)(cmd, namespace)
 
@@ -1541,6 +1577,30 @@ def process_nw_flow_log_create_namespace(cmd, namespace):
 
 def process_nw_flow_log_set_namespace(cmd, namespace):
     from msrestazure.tools import is_valid_resource_id, resource_id
+    if namespace.vnet and not is_valid_resource_id(namespace.vnet):
+        namespace.vnet = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx),
+            resource_group=namespace.resource_group_name,
+            namespace='Microsoft.Network',
+            type='virtualNetworks',
+            name=namespace.vnet)
+
+    if namespace.subnet and not is_valid_resource_id(namespace.subnet):
+        namespace.subnet = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx),
+            resource_group=namespace.resource_group_name,
+            namespace='Microsoft.Network',
+            type='virtualNetworks',
+            name=namespace.subnet)
+
+    if namespace.nic and not is_valid_resource_id(namespace.nic):
+        namespace.nic = resource_id(
+            subscription=get_subscription_id(cmd.cli_ctx),
+            resource_group=namespace.resource_group_name,
+            namespace='Microsoft.Network',
+            type='networkInterfaces',
+            name=namespace.nic)
+
     if namespace.storage_account and not is_valid_resource_id(namespace.storage_account):
         namespace.storage_account = resource_id(
             subscription=get_subscription_id(cmd.cli_ctx),

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

@@ -1213,7 +1213,7 @@ def _make_singular(value):
 
     with self.command_group('network watcher flow-log', client_factory=cf_network_watcher, min_api='2016-09-01') as g:
         g.custom_command('configure',
-                         'set_nsg_flow_logging',
+                         'set_nw_flow_logging',
                          validator=process_nw_flow_log_set_namespace,
                          deprecate_info=self.deprecate(redirect='network watcher flow-log create', hide=False))
         g.custom_show_command('show', 'show_nsg_flow_logging', validator=process_nw_flow_log_show_namespace)

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

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 from collections import Counter, OrderedDict
 
+
 from msrestazure.tools import parse_resource_id, is_valid_resource_id, resource_id
 
 from knack.log import get_logger
@@ -6038,16 +6039,41 @@ def create_nw_packet_capture(cmd, client, resource_group_name, capture_name, vm,
     return client.begin_create(watcher_rg, watcher_name, capture_name, capture_params)
 
 
-def set_nsg_flow_logging(cmd, client, watcher_rg, watcher_name, nsg, storage_account=None,
+def set_vnet_flow_logging(cmd, client, watcher_rg, watcher_name, vnet, resource_group_name=None):
+    vnet_id = _process_vnet_name_and_id(vnet,cmd,resource_group_name)
+    flowlog_status_parameters = cmd.get_models('FlowLogStatusParameters')(target_resource_id=vnet_id)
+    return flowlog_status_parameters
+
+def set_subnet_flow_logging(cmd, subnet, resource_group_name=None):
+    subnet_id = _process_subnet_name_and_id(subnet,cmd,resource_group_name)
+    flowlog_status_parameters = cmd.get_models('FlowLogStatusParameters')(target_resource_id=subnet_id)
+    return flowlog_status_parameters
+
+def set_nic_flow_logging(cmd, nic):
+    flowlog_status_parameters = cmd.get_models('FlowLogStatusParameters')(target_resource_id=nic)
+    return flowlog_status_parameters
+
+def set_nsg_flow_logging(cmd, nsg):
+    flowlog_status_parameters = cmd.get_models('FlowLogStatusParameters')(target_resource_id=nsg)
+    return flowlog_status_parameters
+
+
+def set_nw_flow_logging(cmd, client, watcher_rg, watcher_name, nsg, vnet=None, subnet=None, nic=None, storage_account=None,
                          resource_group_name=None, enabled=None, retention=0, log_format=None, log_version=None,
                          traffic_analytics_workspace=None, traffic_analytics_interval=None,
                          traffic_analytics_enabled=None):
     from azure.cli.core.commands import LongRunningOperation
-    flowlog_status_parameters = cmd.get_models('FlowLogStatusParameters')(target_resource_id=nsg)
+    flowlog_status_parameters = set_nsg_flow_logging(cmd, nsg)
+
+    if vnet!=None:
+        flowlog_status_parameters = set_vnet_flow_logging(cmd,client,watcher_rg,watcher_name,vnet,resource_group_name)
+    elif subnet!=None:
+        flowlog_status_parameters = set_subnet_flow_logging(cmd,client,watcher_rg,watcher_name,subnet,resource_group_name)
+    elif nic!=None:
+        flowlog_status_parameters = set_nic_flow_logging(cmd,client,watcher_rg,watcher_name,nic,resource_group_name)
     config = LongRunningOperation(cmd.cli_ctx)(client.begin_get_flow_log_status(watcher_rg,
                                                                                 watcher_name,
                                                                                 flowlog_status_parameters))
-
     try:
         if not config.flow_analytics_configuration.network_watcher_flow_analytics_configuration.workspace_id:
             config.flow_analytics_configuration = None
@@ -6122,6 +6148,42 @@ def show_nsg_flow_logging(cmd, client, watcher_rg, watcher_name, location=None,
     client = cf_flow_logs(cmd.cli_ctx, None)
     return client.get(watcher_rg, watcher_name, flow_log_name)
 
+def create_vnet_flow_log(cmd, location, vnet, storage_account, resource_group_name, enabled, tags):
+    vnet_id=_process_vnet_name_and_id(vnet,cmd,resource_group_name)
+    FlowLog = cmd.get_models('FlowLog')
+    flow_log = FlowLog(location=location,
+                       target_resource_id=vnet_id,
+                       storage_id=storage_account,
+                       enabled=enabled,
+                       tags=tags)
+    return flow_log
+
+def create_subnet_flow_log(cmd, location, subnet, storage_account, enabled, tags):
+    FlowLog = cmd.get_models('FlowLog')
+    flow_log = FlowLog(location=location,
+                       target_resource_id=subnet,
+                       storage_id=storage_account,
+                       enabled=enabled,
+                       tags=tags)
+    return flow_log
+
+def create_nic_flow_log(cmd, location, nic, storage_account, enabled, tags):
+    FlowLog = cmd.get_models('FlowLog')
+    flow_log = FlowLog(location=location,
+                       target_resource_id=nic,
+                       storage_id=storage_account,
+                       enabled=enabled,
+                       tags=tags)
+    return flow_log
+
+def create_nsg_flow_log(cmd, location, nsg, storage_account, enabled, tags):
+    FlowLog = cmd.get_models('FlowLog')
+    flow_log = FlowLog(location=location,
+                       target_resource_id=nsg,
+                       storage_id=storage_account,
+                       enabled=enabled,
+                       tags=tags)
+    return flow_log
 
 def create_nw_flow_log(cmd,
                        client,
@@ -6130,6 +6192,9 @@ def create_nw_flow_log(cmd,
                        watcher_name,
                        flow_log_name,
                        nsg,
+                       vnet=None, 
+                       subnet=None, 
+                       nic=None,
                        storage_account=None,
                        resource_group_name=None,
                        enabled=None,
@@ -6140,12 +6205,14 @@ def create_nw_flow_log(cmd,
                        traffic_analytics_interval=60,
                        traffic_analytics_enabled=None,
                        tags=None):
-    FlowLog = cmd.get_models('FlowLog')
-    flow_log = FlowLog(location=location,
-                       target_resource_id=nsg,
-                       storage_id=storage_account,
-                       enabled=enabled,
-                       tags=tags)
+    flow_log = create_nsg_flow_log(cmd, location, nsg, storage_account, enabled, tags)
+
+    if vnet != None:
+        flow_log = create_vnet_flow_log(cmd, location, vnet, storage_account, resource_group_name, enabled, tags)
+    elif subnet != None:
+        flow_log = create_subnet_flow_log(cmd,client,watcher_rg,watcher_name,subnet,resource_group_name)
+    elif nic != None:
+        flow_log = create_nic_flow_log(cmd,client,watcher_rg,watcher_name,nic,resource_group_name)
 
     if retention > 0:
         RetentionPolicyParameters = cmd.get_models('RetentionPolicyParameters')
@@ -6196,6 +6263,9 @@ def update_nw_flow_log(cmd,
                        resource_group_name=None,    # dummy parameter to let it appear in command
                        enabled=None,
                        nsg=None,
+                       vnet=None, 
+                       subnet=None, 
+                       nic=None,
                        storage_account=None,
                        retention=0,
                        log_format=None,
@@ -6208,7 +6278,14 @@ def update_nw_flow_log(cmd,
         c.set_param('enabled', enabled)
         c.set_param('tags', tags)
         c.set_param('storage_id', storage_account)
-        c.set_param('target_resource_id', nsg)
+        if vnet!=None:
+            c.set_param('target_resource_id', vnet)
+        elif subnet!=None:
+            c.set_param('target_resource_id', subnet)
+        elif nic!=None:
+            c.set_param('target_resource_id', nic)
+        else:
+            c.set_param('target_resource_id', nsg)
 
     with cmd.update_context(instance.retention_policy) as c:
         c.set_param('days', retention)

src/azure-cli/azure/cli/command_modules/network/tests/latest/test_nw_flow_log.py

@@ -17,6 +17,9 @@ def test_nw_flow_log_create(self, resource_group, resource_group_location, stora
             'location': resource_group_location,
             'storage_account': storage_account,
             'nsg': 'nsg1',
+            'vnet': 'vnet1',
+            'subnet': 'subnet1',
+            'nic': 'nic1',
             'watcher_rg': 'NetworkWatcherRG',
             'watcher_name': 'NetworkWatcher_{}'.format(resource_group_location),
             'flow_log': 'flow_log_test',
@@ -42,6 +45,9 @@ def test_nw_flow_log_create(self, resource_group, resource_group_location, stora
                  '--location {location} '
                  '--resource-group {rg} '
                  '--nsg {nsg} '
+                 '--vnet {vnet} '
+                 '--subnet {subnet} '
+                 '--nic {nic} '
                  '--storage-account {storage_account} '
                  '--workspace {workspace_id} '
                  '--name {flow_log} ')
@@ -67,6 +73,9 @@ def test_nw_flow_log_delete(self, resource_group, resource_group_location, stora
             'location': resource_group_location,
             'storage_account': storage_account,
             'nsg': 'nsg1',
+            'vnet': 'vnet1',
+            'subnet': 'subnet1',
+            'nic': 'nic1',
             'watcher_rg': 'NetworkWatcherRG',
             'watcher_name': 'NetworkWatcher_{}'.format(resource_group_location),
             'flow_log': 'flow_log_test2',
@@ -185,6 +194,9 @@ def test_nw_flow_log_update(self, resource_group, resource_group_location, stora
             'storage_account': storage_account,
             'storage_account_2': 'storageaccount0395',
             'nsg': 'nsg1',
+            'vnet': 'vnet1',
+            'subnet': 'subnet1',
+            'nic': 'nic1',
             'watcher_rg': 'NetworkWatcherRG',
             'watcher_name': 'NetworkWatcher_{}'.format(resource_group_location),
             'flow_log': 'flow_log_test2',