[vm-repair] Bug fix for disk copy V2 to V1 and Adding preview parameter

src/vm-repair/HISTORY.rst

@@ -2,6 +2,14 @@
 Release History
 ===============
 
+0.3.9
+++++++
+Add support for preview flag and fix Gen2 bug
+
+0.3.8
+++++++
+Add support for optional public IP 
+
 0.3.6
 ++++++
 Add support for ALAR2 which requires cloud-init script to prepare the recovery VM with a

src/vm-repair/azext_vm_repair/_help.py

@@ -52,6 +52,9 @@
         - name: Run a local custom script on the VM.
           text: >
             az vm repair run -g MyResourceGroup -n MySourceWinVM --custom-script-file ./file.ps1 --verbose
+        - name: Run unverified script from your fork of https://github.com/Azure/repair-script-library
+          text: >
+            az vm repair run -g MyResourceGroup -n MySourceWinVM --preview "https://github.com/haagha/repair-script-library/blob/master/map.json" --run-id test
 """
 
 helps['vm repair list-scripts'] = """
@@ -67,4 +70,7 @@
         - name: List scripts with test in its description.
           text: >
             az vm repair list-scripts --query "[?contains(description, 'test')]"
+        - name: List unverified script from your fork of https://github.com/Azure/repair-script-library
+          text: >
+            az vm repair list-scripts --preview "https://github.com/haagha/repair-script-library/blob/master/map.json"
 """

src/vm-repair/azext_vm_repair/_params.py

@@ -43,3 +43,7 @@ def load_arguments(self, _):
         c.argument('custom_script_file', help='Custom script file to run on VM. Script should be PowerShell for windows, Bash for Linux.')
         c.argument('parameters', nargs='+', help="Space-separated parameters in the format of '[name=]value'. Positional for bash scripts.")
         c.argument('run_on_repair', help="Script will be run on the linked repair VM.")
+        c.argument('preview', help="URL of forked repair script library's map.json https://github.com/{user}/repair-script-library/blob/master/map.json")
+
+    with self.argument_context('vm repair list-scripts') as c:
+        c.argument('preview', help="URL of forked repair script library's map.json https://github.com/{user}/repair-script-library/blob/master/map.json")

src/vm-repair/azext_vm_repair/custom.py

@@ -12,6 +12,7 @@
 from azure.cli.command_modules.vm.custom import get_vm, _is_linux_os
 from azure.cli.command_modules.storage.storage_url_helpers import StorageResourceIdentifier
 from msrestazure.tools import parse_resource_id
+from .exceptions import SkuDoesNotSupportHyperV
 
 from .command_helper_class import command_helper
 from .repair_utils import (
@@ -32,20 +33,27 @@
     _unlock_singlepass_encrypted_disk,
     _invoke_run_command,
     _check_hyperV_gen,
-    _get_cloud_init_script
+    _get_cloud_init_script,
+    _set_repair_map_url,
+    _is_gen2
 )
 from .exceptions import AzCommandError, SkuNotAvailableError, UnmanagedDiskCopyError, WindowsOsNotAvailableError, RunScriptNotFoundForIdError, SkuDoesNotSupportHyperV, ScriptReturnsError
 logger = get_logger(__name__)
 
 
 def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False):
+
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair create')
     # Main command calling block
     try:
         # Fetch source VM data
         source_vm = get_vm(cmd, resource_group_name, vm_name)
+        source_vm_instance_view = get_vm(cmd, resource_group_name, vm_name, 'instanceView')
+
         is_linux = _is_linux_os(source_vm)
+        is_gen2 = _is_gen2(source_vm_instance_view)
+
         target_disk_name = source_vm.storage_profile.os_disk.name
         is_managed = _uses_managed_disk(source_vm)
         copy_disk_id = None
@@ -59,9 +67,6 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
         else:
             os_image_urn = _fetch_compatible_windows_os_urn(source_vm)
             os_type = 'Windows'
-        # check hyperv Generation
-        if enable_nested:
-            _check_hyperV_gen(source_vm)
 
         # Set up base create vm command
         if is_linux:
@@ -96,6 +101,7 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
             disk_sku, location, os_type, hyperV_generation = _fetch_disk_info(resource_group_name, target_disk_name)
             copy_disk_command = 'az disk create -g {g} -n {n} --source {s} --sku {sku} --location {loc} --os-type {os_type} --query id -o tsv' \
                                 .format(g=resource_group_name, n=copy_disk_name, s=target_disk_name, sku=disk_sku, loc=location, os_type=os_type)
+
             # Only add hyperV variable when available
             if hyperV_generation:
                 copy_disk_command += ' --hyper-v-generation {hyperV}'.format(hyperV=hyperV_generation)
@@ -177,8 +183,8 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
         if enable_nested:
             logger.info("Running Script win-enable-nested-hyperv.ps1 to install HyperV")
 
-            run_hyperv_command = "az vm repair run -g {g} -n {name} --run-id win-enable-nested-hyperv" \
-                .format(g=repair_group_name, name=repair_vm_name)
+            run_hyperv_command = "az vm repair run -g {g} -n {name} --run-id win-enable-nested-hyperv --parameters gen={gen}" \
+                .format(g=repair_group_name, name=repair_vm_name, gen=is_gen2)
             ret_enable_nested = _call_az_command(run_hyperv_command)
 
             logger.debug("az vm repair run hyperv command returned: %s", ret_enable_nested)
@@ -343,12 +349,14 @@ def restore(cmd, vm_name, resource_group_name, disk_name=None, repair_vm_id=None
     return return_dict
 
 
-def run(cmd, vm_name, resource_group_name, run_id=None, repair_vm_id=None, custom_script_file=None, parameters=None, run_on_repair=False):
+def run(cmd, vm_name, resource_group_name, run_id=None, repair_vm_id=None, custom_script_file=None, parameters=None, run_on_repair=False, preview=None):
 
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair run')
     LINUX_RUN_SCRIPT_NAME = 'linux-run-driver.sh'
     WINDOWS_RUN_SCRIPT_NAME = 'win-run-driver.ps1'
+    if preview:
+        _set_repair_map_url(preview)
 
     try:
         # Fetch VM data
@@ -373,6 +381,9 @@ def run(cmd, vm_name, resource_group_name, run_id=None, repair_vm_id=None, custo
             # Fetch run path from GitHub
             repair_script_path = _fetch_run_script_path(run_id)
             run_command_params.append('script_path="./{}"'.format(repair_script_path))
+
+            if preview:
+                run_command_params.append('preview_path="{}"'.format(preview))
         # Custom script scenario for script testers
         else:
             run_command_params.append('script_path=no-op')
@@ -469,10 +480,12 @@ def run(cmd, vm_name, resource_group_name, run_id=None, repair_vm_id=None, custo
     return return_dict
 
 
-def list_scripts(cmd):
+def list_scripts(cmd, preview=None):
 
     # Init command helper object
     command = command_helper(logger, cmd, 'vm repair list-scripts')
+    if preview:
+        _set_repair_map_url(preview)
 
     try:
         run_map = _fetch_run_script_map()

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -35,12 +35,30 @@ def _get_cloud_init_script():
     return os.path.join(rootpath, SCRIPTS_DIR_NAME, CLOUD_INIT)
 
 
+def _set_repair_map_url(url):
+    raw_url = str(url)
+    if "github.com" in raw_url:
+        raw_url = raw_url.replace("github.com", "raw.githubusercontent.com")
+        raw_url = raw_url.replace("/blob/", "/")
+        global REPAIR_MAP_URL
+        REPAIR_MAP_URL = raw_url
+        print(REPAIR_MAP_URL)
+
+
 def _uses_managed_disk(vm):
     if vm.storage_profile.os_disk.managed_disk is None:
         return False
     return True
 
 
+def _is_gen2(vm):
+    gen = 1
+    gen = vm.instance_view.hyper_v_generation
+    if gen.lower() == 'v2':
+        return 2
+    return 1
+
+
 def _call_az_command(command_string, run_async=False, secure_params=None):
     """
     Uses subprocess to run a command string. To hide sensitive parameters from logs, add the
@@ -237,7 +255,7 @@ def _fetch_compatible_sku(source_vm, hyperv):
 
 def _fetch_disk_info(resource_group_name, disk_name):
     """ Returns sku, location, os_type, hyperVgeneration as tuples """
-    show_disk_command = 'az disk show -g {g} -n {name} --query [sku.name,location,osType,hyperVgeneration] -o json'.format(g=resource_group_name, name=disk_name)
+    show_disk_command = 'az disk show -g {g} -n {name} --query [sku.name,location,osType,hyperVGeneration] -o json'.format(g=resource_group_name, name=disk_name)
     disk_info = loads(_call_az_command(show_disk_command))
     # Note that disk_info will always have 4 elements if the command succeeded, if it fails it will cause an exception
     sku, location, os_type, hyper_v_version = disk_info[0], disk_info[1], disk_info[2], disk_info[3]

src/vm-repair/azext_vm_repair/tests/latest/test_repair_commands.py

@@ -24,7 +24,7 @@ def test_vmrepair_WinManagedCreateRestore(self, resource_group):
         assert len(vms) == 1
 
         # Test create
-        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2018 -o json').get_output_in_json()
+        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2018 --associate-public-ip -o json').get_output_in_json()
         assert result['status'] == STATUS_SUCCESS, result['error_message']
 
         # Check repair VM
@@ -521,3 +521,37 @@ def test_vmrepair_LinuxRunHelloWorld(self, resource_group):
 
         # Check Output
         assert 'Hello World!' in result['output']
+
+
+class WindowsManagedDiskCreateRestoreGen2Test(LiveScenarioTest):
+
+    @ResourceGroupPreparer(location='westus2')
+    def test_vmrepair_WinManagedCreateRestore(self, resource_group):
+        self.kwargs.update({
+            'vm': 'vm1'
+        })
+
+        # Create test VM
+        self.cmd('vm create -g {rg} -n {vm} --admin-username azureadmin --image MicrosoftWindowsServer:windowsserver-gen2preview:2019-datacenter-gen2:2019.0.20190620 --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
+
+        # Test create
+        result = self.cmd('vm repair create -g {rg} -n {vm} --repair-username azureadmin --repair-password !Passw0rd2018 --associate-public-ip -o json').get_output_in_json()
+        assert result['status'] == STATUS_SUCCESS, result['error_message']
+
+        # Check repair VM
+        repair_vms = self.cmd('vm list -g {} -o json'.format(result['repair_resource_group'])).get_output_in_json()
+        assert len(repair_vms) == 1
+        repair_vm = repair_vms[0]
+        # Check attached data disk
+        assert repair_vm['storageProfile']['dataDisks'][0]['name'] == result['copied_disk_name']
+
+        # Call Restore
+        self.cmd('vm repair restore -g {rg} -n {vm} --yes')
+
+        # Check swapped OS disk
+        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']

src/vm-repair/setup.py

@@ -8,7 +8,7 @@
 from codecs import open
 from setuptools import setup, find_packages
 
-VERSION = "0.3.8"
+VERSION = "0.3.9"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',