Got a solution from MS:
https://github.com/Azure/azure-sdk-for-python/issues/35598
For linux:
from azure.mgmt.compute.models import VirtualMachineExtension
response = client.virtual_machine_extensions.begin_create_or_update(
resource_group_name={resource_group_name},
vm_name={vm_name},
vm_extension_name='enablevmaccess',
extension_parameters=VirtualMachineExtension(
location={location},
protected_settings={'username': {username}, 'password': {password}}))
For Windows:
response = client.virtual_machine_extensions.begin_create_or_update(
resource_group_name={resource_group_name},
vm_name={vm_name},
vm_extension_name='enablevmaccess',
extension_parameters=VirtualMachineExtension(
location={location},
publisher='Microsoft.Compute',
type_properties_type='VMAccessAgent',
type_handler_version='2.4',
protected_settings={'password': {password}},
settings={'userName': {username}}))