associate subnets with NSG (#1393)

* associate subnets with NSG

change NSG rule protocol to ANY

* subnet wait

* Improve NSG update logic

validate that subnet nsg is set before getting it's id (#1409)

Co-authored-by: stas <statis@microsoft.com>

src/api-service/__app__/onefuzzlib/azure/subnet.py

@@ -8,6 +8,7 @@ import os
 from typing import Any, Optional, Union, cast
 
 from azure.core.exceptions import ResourceNotFoundError
+from azure.mgmt.network.models import Subnet, VirtualNetwork
 from msrestazure.azure_exceptions import CloudError
 from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error, NetworkConfig
@@ -16,21 +17,41 @@ from onefuzztypes.primitives import Region
 from .network_mgmt_client import get_network_client
 
 
-def get_subnet_id(resource_group: str, name: str, subnet_name: str) -> Optional[str]:
+def get_vnet(resource_group: str, name: str) -> Optional[VirtualNetwork]:
     network_client = get_network_client()
     try:
-        subnet = network_client.subnets.get(resource_group, name, subnet_name)
-        return cast(str, subnet.id)
+        vnet = network_client.virtual_networks.get(resource_group, name)
+        return cast(VirtualNetwork, vnet)
     except (CloudError, ResourceNotFoundError):
         logging.info(
-            "subnet missing: resource group:%s name:%s subnet_name:%s",
+            "vnet missing: resource group:%s name:%s",
             resource_group,
             name,
-            subnet_name,
         )
     return None
 
 
+def get_subnet(
+    resource_group: str, vnet_name: str, subnet_name: str
+) -> Optional[Subnet]:
+    # Has to get using vnet. That way NSG field is properly set up in subnet
+    vnet = get_vnet(resource_group, vnet_name)
+    if vnet:
+        for subnet in vnet.subnets:
+            if subnet.name == subnet_name:
+                return subnet
+
+    return None
+
+
+def get_subnet_id(resource_group: str, name: str, subnet_name: str) -> Optional[str]:
+    subnet = get_subnet(resource_group, name, subnet_name)
+    if subnet:
+        return cast(str, subnet.id)
+    else:
+        return None
+
+
 def delete_subnet(resource_group: str, name: str) -> Union[None, CloudError, Any]:
     network_client = get_network_client()
     try: