use standard pip and add nsg rule

LiliDeng · LiliDeng · commit 917a4e2bea5b · 2025-04-23T16:12:34.000+08:00
diff --git a/lisa/sut_orchestrator/azure/arm_template.bicep b/lisa/sut_orchestrator/azure/arm_template.bicep
@@ -55,6 +55,9 @@ param use_ipv6 bool = false
 @description('whether to enable network outbound access')
 param enable_vm_nat bool
 
+@description('The source IP address prefixes allowed in NSG')
+param source_address_prefixes array
+
 var vnet_id = virtual_network_name_resource.id
 var node_count = length(nodes)
 var availability_set_name_value = 'lisa-availabilitySet'
@@ -253,9 +256,65 @@ resource virtual_network_name_resource 'Microsoft.Network/virtualNetworks@2024-0
           use_ipv6 ? ['2001:db8:${j}::/64'] : []
         )
         defaultOutboundAccess: enable_vm_nat
+        networkSecurityGroup: {
+          id: resourceId('Microsoft.Network/networkSecurityGroups', '${toLower(virtual_network_name)}-nsg')
+        }
       }
     }]
   }
+  dependsOn: [
+    nsg
+  ]
+}
+
+resource nsg 'Microsoft.Network/networkSecurityGroups@2024-05-01' = {
+  name: '${toLower(virtual_network_name)}-nsg'
+  location: location
+  properties: {
+    securityRules: [
+      {
+        name: 'LISASSH'
+        properties: {
+          priority: 100
+          direction: 'Inbound'
+          access: 'Allow'
+          protocol: 'Tcp'
+          sourcePortRange: '*'
+          destinationPortRange: '22'
+          sourceAddressPrefixes: source_address_prefixes
+          destinationAddressPrefix: '*'
+        }
+      }
+      {
+          name: 'LISAKVMSSH'
+          properties: {
+              description: 'Allows nested VM SSH traffic'
+              protocol: 'Tcp'
+              sourcePortRange: '*'
+              destinationPortRange: '60020-60030'
+              destinationAddressPrefix: '*'
+              sourceAddressPrefixes: source_address_prefixes
+              access: 'Allow'
+              priority: 206
+              direction: 'Inbound'
+          }
+      }
+      {
+        name: 'LISALIBVIRTSSH'
+        properties: {
+            description: 'Allows SSH traffic to Libvirt Platform Guests'
+            protocol: 'Tcp'
+            sourcePortRange: '*'
+            destinationPortRange: '49152-49352'
+            destinationAddressPrefix: '*'
+            sourceAddressPrefixes: source_address_prefixes
+            access: 'Allow'
+            priority: 208
+            direction: 'Inbound'
+        }
+      }
+    ]
+  }
 }
 
 resource availability_set 'Microsoft.Compute/availabilitySets@2019-07-01' = if (use_availability_set) {
@@ -273,11 +332,11 @@ resource nodes_public_ip 'Microsoft.Network/publicIPAddresses@2020-05-01' = [for
   tags: tags
   name: '${nodes[i].name}-public-ip'
   properties: {
-    publicIPAllocationMethod: ((is_ultradisk || use_availability_zones || use_ipv6) ? 'Static' : 'Dynamic')
+    publicIPAllocationMethod: 'Static'
     ipTags: (empty(ip_tags) ? null : ip_tags)
   }
   sku: {
-    name: ((is_ultradisk || use_availability_zones || use_ipv6) ? 'Standard' : 'Basic')
+    name: 'Standard'
   }
   zones: (use_availability_zones ? availability_zones : null)
 }]
diff --git a/lisa/sut_orchestrator/azure/autogen_arm_template.json b/lisa/sut_orchestrator/azure/autogen_arm_template.json
@@ -6,7 +6,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.33.93.31351",
-      "templateHash": "15877188685327638829"
+      "templateHash": "14741729812904551283"
     }
   },
   "functions": [
@@ -530,6 +530,12 @@
       "metadata": {
         "description": "whether to enable network outbound access"
       }
+    },
+    "source_address_prefixes": {
+      "type": "array",
+      "metadata": {
+        "description": "The source IP address prefixes allowed in NSG"
+      }
     }
   },
   "variables": {
@@ -574,14 +580,71 @@
               "name": "[format('{0}{1}', parameters('subnet_prefix'), range(0, parameters('subnet_count'))[copyIndex('subnets')])]",
               "properties": {
                 "addressPrefixes": "[concat(createArray(format('10.0.{0}.0/24', range(0, parameters('subnet_count'))[copyIndex('subnets')])), if(parameters('use_ipv6'), createArray(format('2001:db8:{0}::/64', range(0, parameters('subnet_count'))[copyIndex('subnets')])), createArray()))]",
-                "defaultOutboundAccess": "[parameters('enable_vm_nat')]"
+                "defaultOutboundAccess": "[parameters('enable_vm_nat')]",
+                "networkSecurityGroup": {
+                  "id": "[resourceId('Microsoft.Network/networkSecurityGroups', format('{0}-nsg', toLower(parameters('virtual_network_name'))))]"
+                }
               }
             }
           }
         ],
         "addressSpace": {
           "addressPrefixes": "[concat(createArray('10.0.0.0/16'), if(parameters('use_ipv6'), createArray('2001:db8::/32'), createArray()))]"
         }
+      },
+      "dependsOn": [
+        "nsg"
+      ]
+    },
+    "nsg": {
+      "type": "Microsoft.Network/networkSecurityGroups",
+      "apiVersion": "2024-05-01",
+      "name": "[format('{0}-nsg', toLower(parameters('virtual_network_name')))]",
+      "location": "[parameters('location')]",
+      "properties": {
+        "securityRules": [
+          {
+            "name": "LISASSH",
+            "properties": {
+              "priority": 100,
+              "direction": "Inbound",
+              "access": "Allow",
+              "protocol": "Tcp",
+              "sourcePortRange": "*",
+              "destinationPortRange": "22",
+              "sourceAddressPrefixes": "[parameters('source_address_prefixes')]",
+              "destinationAddressPrefix": "*"
+            }
+          },
+          {
+            "name": "LISAKVMSSH",
+            "properties": {
+              "description": "Allows nested VM SSH traffic",
+              "protocol": "Tcp",
+              "sourcePortRange": "*",
+              "destinationPortRange": "60020-60030",
+              "destinationAddressPrefix": "*",
+              "sourceAddressPrefixes": "[parameters('source_address_prefixes')]",
+              "access": "Allow",
+              "priority": 206,
+              "direction": "Inbound"
+            }
+          },
+          {
+            "name": "LISALIBVIRTSSH",
+            "properties": {
+              "description": "Allows SSH traffic to Libvirt Platform Guests",
+              "protocol": "Tcp",
+              "sourcePortRange": "*",
+              "destinationPortRange": "49152-49352",
+              "destinationAddressPrefix": "*",
+              "sourceAddressPrefixes": "[parameters('source_address_prefixes')]",
+              "access": "Allow",
+              "priority": 208,
+              "direction": "Inbound"
+            }
+          }
+        ]
       }
     },
     "availability_set": {
@@ -607,11 +670,11 @@
       "location": "[parameters('location')]",
       "tags": "[parameters('tags')]",
       "properties": {
-        "publicIPAllocationMethod": "[if(or(or(parameters('is_ultradisk'), variables('use_availability_zones')), parameters('use_ipv6')), 'Static', 'Dynamic')]",
+        "publicIPAllocationMethod": "Static",
         "ipTags": "[if(empty(variables('ip_tags')), null(), variables('ip_tags'))]"
       },
       "sku": {
-        "name": "[if(or(or(parameters('is_ultradisk'), variables('use_availability_zones')), parameters('use_ipv6')), 'Standard', 'Basic')]"
+        "name": "Standard"
       },
       "zones": "[if(variables('use_availability_zones'), variables('availability_zones'), null())]"
     },
diff --git a/lisa/sut_orchestrator/azure/common.py b/lisa/sut_orchestrator/azure/common.py
@@ -1209,6 +1209,7 @@ class AzureArmParameter:
     is_ultradisk: bool = False
     use_ipv6: bool = False
     enable_vm_nat: bool = False
+    source_address_prefixes: List[str] = field(default_factory=list)
 
     def __post_init__(self, *args: Any, **kwargs: Any) -> None:
         add_secret(self.admin_username, PATTERN_HEADTAIL)
diff --git a/lisa/sut_orchestrator/azure/platform_.py b/lisa/sut_orchestrator/azure/platform_.py
@@ -80,6 +80,7 @@
     get_first_combination,
     get_matched_str,
     get_or_generate_key_pairs,
+    get_public_ip,
     get_public_key_data,
     is_unittest,
     plugin_manager,
@@ -301,6 +302,7 @@ class AzurePlatformSchema:
     # will be retired. It's recommended to disable outbound access to
     # enforce explicit connectivity rules.
     enable_vm_nat: bool = field(default=False)
+    source_address_prefixes: Optional[List[str]] = field(default=None)
 
     virtual_network_resource_group: str = field(default="")
     virtual_network_name: str = field(default=AZURE_VIRTUAL_NETWORK_NAME)
@@ -356,6 +358,7 @@ def __post_init__(self, *args: Any, **kwargs: Any) -> None:
                 "use_public_address",
                 "use_ipv6",
                 "enable_vm_nat",
+                "source_address_prefixes",
             ],
         )
 
@@ -939,6 +942,7 @@ def _initialize(self, *args: Any, **kwargs: Any) -> None:
         self.subscription_id = azure_runbook.subscription_id
         self.cloud = azure_runbook.cloud
         self.resource_group_managed_by = azure_runbook.resource_group_managed_by
+        self._cached_ip_address: List[str] = []
 
         self._initialize_credential()
 
@@ -956,6 +960,15 @@ def _initialize(self, *args: Any, **kwargs: Any) -> None:
             self.credential, self.subscription_id, self.cloud
         )
 
+    def _get_ip_addresses(self) -> List[str]:
+        if self._cached_ip_address:
+            return self._cached_ip_address
+        if self._azure_runbook.source_address_prefixes:
+            self._cached_ip_address = self._azure_runbook.source_address_prefixes
+        else:
+            self._cached_ip_address = [get_public_ip()]
+        return self._cached_ip_address
+
     def _initialize_credential(self) -> None:
         azure_runbook = self._azure_runbook
         if azure_runbook.credential:
@@ -1254,6 +1267,8 @@ def _create_deployment_parameters(
             self._azure_runbook.shared_resource_group_name
         )
         arm_parameters.enable_vm_nat = self._azure_runbook.enable_vm_nat
+        arm_parameters.source_address_prefixes = self._get_ip_addresses()
+
         # the arm template may be updated by the hooks, so make a copy to avoid
         # the original template is modified.
         template = deepcopy(self._load_template())
diff --git a/lisa/util/__init__.py b/lisa/util/__init__.py
@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
+import ipaddress
 import random
 import re
 import string
@@ -28,9 +29,11 @@
 
 import paramiko
 import pluggy
+import requests
 from assertpy import assert_that
 from dataclasses_json import config
 from marshmallow import fields
+from retry import retry
 from semver import VersionInfo
 
 from lisa import secret
@@ -947,3 +950,11 @@ def to_bool(value: Union[str, bool, int]) -> bool:
     raise TypeError(
         f"Unsupported type for conversion to boolean: {type(value).__name__}"
     )
+
+
+@retry(tries=10, delay=0.5)
+def get_public_ip() -> str:
+    response = requests.get("https://api.ipify.org/", timeout=5)
+    result = response.text
+    ipaddress.ip_address(result)
+    return str(result)
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,6 +27,7 @@ dependencies = [
     "spurplus ~= 2.3.5",
     "websockets ~= 10.3",
     "charset_normalizer ~= 2.1.1",
+    "requests ~= 2.32.0",
 ]
 dynamic = ["version"]
 license = {text = "MIT"}
@@ -58,13 +59,11 @@ azure = [
     "cachetools ~= 5.2.0",
     "Pillow <= 11.1.0",
     "PyGObject <= 3.50.0; platform_system == 'Linux'",
-    "requests ~= 2.32.0",
     "pycdlib ~= 1.12.0",
 ]
 
 ado = [
     "azure-devops ~= 7.1.0b3",
-    "requests ~= 2.32.0",
 ]
 
 black = [
@@ -105,7 +104,6 @@ baremetal = [
     "pysmb ~= 1.2.9.1",
     "redfish ~= 3.2.1",
     "azure-devops ~= 7.1.0b3",
-    "requests ~= 2.32.0",
 ]
 
 mypy = [
