Add ASG if exist when nic is reset (#7513)

* Add ASG if exist when nic is reset * Revert "Add ASG if exist when nic is reset" This reverts commit 4429db2. * Add back ASG if exist when nic is reset * stringlize applicationSecurityGroups * ASG for revert ip * add test case --------- Co-authored-by: REDMOND\waylandwu <[email protected]>
Azure · Apr 18, 2024 · 2a4117b · 2a4117b
1 parent 4ea607e
commit 2a4117b
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 8 deletions.
diff --git a/src/vm-repair/HISTORY.rst b/src/vm-repair/HISTORY.rst
@@ -1,6 +1,11 @@
 
 Release History
 ===============
+1.0.5
+++++++
+Bug fix ASG is not added properly when reset the nic
+Add ASG if exist when nic is reset 
+
 1.0.4
 ++++++
 Logging improvements and script fixing

diff --git a/src/vm-repair/azext_vm_repair/custom.py b/src/vm-repair/azext_vm_repair/custom.py
@@ -612,6 +612,16 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
         vnet_resource_group = subnet_id_tokens[-7]
         ipconfig_name = ip_config_object['name']
         orig_ip_address = ip_config_object['privateIPAddress']
+        application_names=""
+        applicationSecurityGroups='applicationSecurityGroups'
+        if applicationSecurityGroups in ip_config_object:
+            for item in ip_config_object[applicationSecurityGroups]:
+                application_id_tokens = item['id'].split('/')
+                if application_id_tokens[-1] is not None:
+                    application_names+=application_id_tokens[-1]+ " "
+            logger.info('applicationSecurityGroups {application_names}...\n')    
+
+
         # Dynamic | Static
         orig_ip_allocation_method = ip_config_object['privateIPAllocationMethod']
 
@@ -626,8 +636,13 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
         # 3) Update private IP address to another in subnet. This will invoke and wait for a VM restart.
         logger.info('Updating VM IP configuration. This might take a few minutes...\n')
         # Update IP address
-        update_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} ' \
-                            .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=swap_ip_address)
+        if application_names:
+            update_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} --asgs {asgs}' \
+                                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=swap_ip_address,asgs=application_names)
+        else:
+            logger.info('applicationSecurityGroups do not exist...\n')
+            update_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip}' \
+                                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=swap_ip_address)
         _call_az_command(update_ip_command)
 
         # 4) Change things back. This will also invoke and wait for a VM restart.
@@ -636,12 +651,20 @@ def reset_nic(cmd, vm_name, resource_group_name, yes=False):
         revert_ip_command = None
         if orig_ip_allocation_method == DYNAMIC_CONFIG:
             # Revert Static to Dynamic
-            revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --set privateIpAllocationMethod={method}' \
-                                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, method=DYNAMIC_CONFIG)
+            if application_names:
+                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --set privateIpAllocationMethod={method} --asgs {asgs}' \
+                                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, method=DYNAMIC_CONFIG,asgs=application_names)
+            else:
+                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --set privateIpAllocationMethod={method}' \
+                                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, method=DYNAMIC_CONFIG)
         else:
             # Revert to original static ip
-            revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} ' \
-                                .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=orig_ip_address)
+            if application_names:
+                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} --asgs {asgs}' \
+                                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=orig_ip_address,asgs=application_names)
+            else:
+                revert_ip_command = 'az network nic ip-config update -g {g} --nic-name {nic} -n {config} --private-ip-address {ip} ' \
+                                    .format(g=resource_group_name, nic=primary_nic_name, config=ipconfig_name, ip=orig_ip_address)
 
         _call_az_command(revert_ip_command)
         logger.info('VM guest NIC reset is complete and all configurations are reverted.')
@@ -762,4 +785,4 @@ def repair_and_restore(cmd, vm_name, resource_group_name, repair_password=None,
 
     logger.info('\n%s\n', command.message)
 
-    return return_dict
+    return return_dict
diff --git a/src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py b/src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py
@@ -781,3 +781,31 @@ def test_vmrepair_LinuxManagedCreateRestore(self, resource_group):
         vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
         source_vm = vms[0]
         assert source_vm['storageProfile']['osDisk']['name'] == result['copied_disk_name']
+
+class ResetNICWithASG(LiveScenarioTest):
+
+    @ResourceGroupPreparer(location='westus2')
+    def test_vmrepair_ResetNicWindowsVM(self, resource_group):
+        self.kwargs.update({
+            'vm': 'vm1'
+        })
+
+        # Create test VM
+        self.cmd('vm create -g {rg} -n {vm} --admin-username azureadmin --image Win2016Datacenter --admin-password !Passw0rd2018')
+        vms = self.cmd('vm list -g {rg} -o json').get_output_in_json()
+        # Something wrong with vm create command if it fails here
+        assert len(vms) == 1
+
+        #create ASG
+        self.cmd('az network asg create \
+                  --resource-group myResourceGroup \
+                  --name myASG')
+
+        # Test Reset NIC
+        self.cmd('vm repair reset-nic -g {rg} -n {vm} --yes')
+
+        # Mac address should be changed in the Guest OS but no way to assert from here.
+        # Assert that the VM is still running afterwards
+        vm_instance_view = self.cmd('vm get-instance-view -g {rg} -n {vm} -o json').get_output_in_json()
+        vm_power_state = vm_instance_view['instanceView']['statuses'][1]['code']
+        assert vm_power_state == 'PowerState/running'
diff --git a/src/vm-repair/setup.py b/src/vm-repair/setup.py
@@ -8,7 +8,7 @@
 from codecs import open
 from setuptools import setup, find_packages
 
-VERSION = "1.0.4"
+VERSION = "1.0.5"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',