From c2b16f2d275c1d5b68bbbf1072b81b1a417dd62e Mon Sep 17 00:00:00 2001
From: Kameron Carr <kameroncarr@microsoft.com>
Date: Tue, 1 Oct 2024 11:32:05 -0700
Subject: [PATCH] ip_service_tags

New security standard will require all IPs to be tagged.

Please note this IP tagging system is different from general  resource tagging.
---
 .../sut_orchestrator/azure/arm_template.bicep |  8 ++++++
 .../azure/autogen_arm_template.json           | 27 +++++++++++++++----
 lisa/sut_orchestrator/azure/common.py         |  1 +
 lisa/sut_orchestrator/azure/platform_.py      |  2 ++
 4 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/lisa/sut_orchestrator/azure/arm_template.bicep b/lisa/sut_orchestrator/azure/arm_template.bicep
index 2366e37aca..d067f9ed3b 100644
--- a/lisa/sut_orchestrator/azure/arm_template.bicep
+++ b/lisa/sut_orchestrator/azure/arm_template.bicep
@@ -46,6 +46,9 @@ param data_disks array
 @description('whether to use ultra disk')
 param is_ultradisk bool = false
 
+@description('IP Service Tags')
+param ip_service_tags object
+
 var vnet_id = virtual_network_name_resource.id
 var node_count = length(nodes)
 var availability_set_name_value = 'lisa-availabilitySet'
@@ -59,6 +62,10 @@ var use_availability_zones = (availability_type == 'availability_zone')
 var availability_set_value = (use_availability_set ? getAvailabilitySetId(availability_set_name_value): null)
 var combined_vm_tags = union(tags, vm_tags)
 var combined_aset_tags = union(tags, availability_set_tags)
+var ip_tags = [for key in objectKeys(ip_service_tags): {
+  ipTagType: key
+  tag: ip_service_tags[key]
+}]
 
 func isCvm(node object) bool => bool((!empty(node.vhd)) && (!empty(node.vhd.vmgs_path)))
 
@@ -254,6 +261,7 @@ resource nodes_public_ip 'Microsoft.Network/publicIPAddresses@2020-05-01' = [for
   name: '${nodes[i].name}-public-ip'
   properties: {
     publicIPAllocationMethod: ((is_ultradisk || use_availability_zones) ? 'Static' : 'Dynamic')
+    ipTags: (empty(ip_tags) ? null : ip_tags)
   }
   sku: {
     name: ((is_ultradisk || use_availability_zones) ? 'Standard' : 'Basic')
diff --git a/lisa/sut_orchestrator/azure/autogen_arm_template.json b/lisa/sut_orchestrator/azure/autogen_arm_template.json
index 0ca54ce928..18c4a3484c 100644
--- a/lisa/sut_orchestrator/azure/autogen_arm_template.json
+++ b/lisa/sut_orchestrator/azure/autogen_arm_template.json
@@ -5,8 +5,8 @@
   "metadata": {
     "_generator": {
       "name": "bicep",
-      "version": "0.28.1.47646",
-      "templateHash": "1630470125196944168"
+      "version": "0.30.23.60470",
+      "templateHash": "17909783643222378721"
     }
   },
   "functions": [
@@ -511,9 +511,25 @@
       "metadata": {
         "description": "whether to use ultra disk"
       }
+    },
+    "ip_service_tags": {
+      "type": "object",
+      "metadata": {
+        "description": "IP Service Tags"
+      }
     }
   },
   "variables": {
+    "copy": [
+      {
+        "name": "ip_tags",
+        "count": "[length(objectKeys(parameters('ip_service_tags')))]",
+        "input": {
+          "ipTagType": "[objectKeys(parameters('ip_service_tags'))[copyIndex('ip_tags')]]",
+          "tag": "[parameters('ip_service_tags')[objectKeys(parameters('ip_service_tags'))[copyIndex('ip_tags')]]]"
+        }
+      }
+    ],
     "vnet_id": "[resourceId('Microsoft.Network/virtualNetworks', parameters('virtual_network_name'))]",
     "node_count": "[length(parameters('nodes'))]",
     "availability_set_name_value": "lisa-availabilitySet",
@@ -579,7 +595,8 @@
       "location": "[parameters('location')]",
       "tags": "[parameters('tags')]",
       "properties": {
-        "publicIPAllocationMethod": "[if(or(parameters('is_ultradisk'), variables('use_availability_zones')), 'Static', 'Dynamic')]"
+        "publicIPAllocationMethod": "[if(or(parameters('is_ultradisk'), variables('use_availability_zones')), 'Static', 'Dynamic')]",
+        "ipTags": "[if(empty(variables('ip_tags')), null(), variables('ip_tags'))]"
       },
       "sku": {
         "name": "[if(or(parameters('is_ultradisk'), variables('use_availability_zones')), 'Standard', 'Basic')]"
@@ -770,8 +787,8 @@
           "metadata": {
             "_generator": {
               "name": "bicep",
-              "version": "0.28.1.47646",
-              "templateHash": "10740733774987815957"
+              "version": "0.30.23.60470",
+              "templateHash": "12249187708601787514"
             }
           },
           "functions": [
diff --git a/lisa/sut_orchestrator/azure/common.py b/lisa/sut_orchestrator/azure/common.py
index 042338eba4..24a2e6974d 100644
--- a/lisa/sut_orchestrator/azure/common.py
+++ b/lisa/sut_orchestrator/azure/common.py
@@ -1154,6 +1154,7 @@ class AzureArmParameter:
     data_disks: List[DataDiskSchema] = field(default_factory=list)
     vm_tags: Dict[str, Any] = field(default_factory=dict)
     tags: Dict[str, Any] = field(default_factory=dict)
+    ip_service_tags: Dict[str, str] = field(default_factory=dict)
 
     virtual_network_resource_group: str = ""
     virtual_network_name: str = AZURE_VIRTUAL_NETWORK_NAME
diff --git a/lisa/sut_orchestrator/azure/platform_.py b/lisa/sut_orchestrator/azure/platform_.py
index 6cc79ce61d..5e58639058 100644
--- a/lisa/sut_orchestrator/azure/platform_.py
+++ b/lisa/sut_orchestrator/azure/platform_.py
@@ -279,6 +279,7 @@ class AzurePlatformSchema:
     vm_tags: Optional[Dict[str, Any]] = field(default=None)
     tags: Optional[Dict[str, Any]] = field(default=None)
     use_public_address: bool = field(default=True)
+    ip_service_tags: Optional[Dict[str, str]] = field(default=None)
 
     virtual_network_resource_group: str = field(default="")
     virtual_network_name: str = field(default=AZURE_VIRTUAL_NETWORK_NAME)
@@ -1058,6 +1059,7 @@ def _create_deployment_parameters(
         copied_fields = [
             "vm_tags",
             "tags",
+            "ip_service_tags",
         ]
         availability_copied_fields = [
             "availability_set_tags",