vm/vmss update migration

Author: naga-nandyala

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

@@ -1649,6 +1649,7 @@ def update_vm(cmd, resource_group_name, vm_name, os_disk=None, disk_caching=None
               security_type=None, enable_proxy_agent=None, proxy_agent_mode=None, additional_scheduled_events=None,
               enable_user_reboot_scheduled_events=None, enable_user_redeploy_scheduled_events=None,
               align_regional_disks_to_vm_zone=None, wire_server_mode=None, imds_mode=None,
+              add_proxy_agent_extension=None,
               wire_server_access_control_profile_reference_id=None, imds_access_control_profile_reference_id=None,
               key_incarnation_id=None, **kwargs):
     from azure.mgmt.core.tools import parse_resource_id, resource_id, is_valid_resource_id
@@ -1680,6 +1681,26 @@ def _output(self, *args, **kwargs):
         }))
         vm = get_vm_to_update(cmd, resource_group_name, vm_name)
 
+    if add_proxy_agent_extension is not None:
+        from .aaz.latest.vm import Update as _Update
+        
+        class Update(_Update):
+            pass
+        
+        args = {
+            'resource_group': resource_group_name,
+            'vm_name': vm_name,
+            'no_wait': no_wait,
+            'security_profile': {
+                'proxy_agent_settings': {
+                    'add_proxy_agent_extension': add_proxy_agent_extension
+                }
+            }
+        }
+        
+        LongRunningOperation(cmd.cli_ctx)(Update(cli_ctx=cmd.cli_ctx)(command_args=args))
+        vm = get_vm_to_update(cmd, resource_group_name, vm_name)
+
     disk_name = None
     if os_disk is not None:
         if is_valid_resource_id(os_disk):
@@ -1897,7 +1918,8 @@ def _output(self, *args, **kwargs):
                 }
     client = _compute_client_factory(cmd.cli_ctx, aux_subscriptions=aux_subscriptions)
     if wire_server_access_control_profile_reference_id is not None or \
-            imds_access_control_profile_reference_id is not None:
+            imds_access_control_profile_reference_id is not None or \
+            add_proxy_agent_extension is not None:
         kwargs['parameters'] = vm
     return sdk_no_wait(no_wait, client.virtual_machines.begin_create_or_update, resource_group_name, vm_name, **kwargs)
 # endregion
@@ -4167,7 +4189,8 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
                 upgrade_policy_mode=None, enable_auto_os_upgrade=None, skuprofile_vmsizes=None,
                 skuprofile_allostrat=None, skuprofile_rank=None,
                 security_posture_reference_is_overridable=None, zone_balance=None,
-                wire_server_mode=None, imds_mode=None, wire_server_access_control_profile_reference_id=None,
+                wire_server_mode=None, imds_mode=None, add_proxy_agent_extension=None,
+                wire_server_access_control_profile_reference_id=None,
                 imds_access_control_profile_reference_id=None, enable_automatic_zone_balancing=None,
                 automatic_zone_balancing_strategy=None, automatic_zone_balancing_behavior=None, **kwargs):
     vmss = kwargs['parameters']
@@ -4198,6 +4221,28 @@ def _output(self, *args, **kwargs):
         }))
         vmss = get_vmss_modified(cmd, resource_group_name, name, instance_id, security_type)
 
+    if add_proxy_agent_extension is not None:
+        from .aaz.latest.vmss import Update as _VMSSUpdate
+        
+        class VMSSUpdate(_VMSSUpdate):
+            pass
+        
+        args = {
+            'resource_group': resource_group_name,
+            'vm_scale_set_name': name,
+            'no_wait': no_wait,
+            'virtual_machine_profile': {
+                'security_profile': {
+                    'proxy_agent_settings': {
+                        'add_proxy_agent_extension': add_proxy_agent_extension
+                    }
+                }
+            }
+        }
+        
+        LongRunningOperation(cmd.cli_ctx)(VMSSUpdate(cli_ctx=cmd.cli_ctx)(command_args=args))
+        vmss = get_vmss_modified(cmd, resource_group_name, name, instance_id, security_type)
+
     aux_subscriptions = None
     # pylint: disable=too-many-boolean-expressions
     if vmss and hasattr(vmss, 'virtual_machine_profile') and vmss.virtual_machine_profile and \
@@ -4543,7 +4588,8 @@ def _output(self, *args, **kwargs):
         vmss.zone_balance = zone_balance
 
     if wire_server_access_control_profile_reference_id is not None or \
-            imds_access_control_profile_reference_id is not None:
+            imds_access_control_profile_reference_id is not None or \
+            add_proxy_agent_extension is not None:
         kwargs['parameters'] = vmss
 
     return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,