vm extension: support 'set' and 'list' (Azure#326)

* list extension image list

* extract common helper logic to vm_utils

.gitignore

@@ -62,3 +62,6 @@ src/build
 
 # Auxiliary files
 command_coverage.txt
+
+#Test artifacts
+private_config.json

azure-cli.pyproj

@@ -291,6 +291,9 @@
     <Compile Include="command_modules\azure-cli-vm\azure\cli\command_modules\vm\_params.py">
       <SubType>Code</SubType>
     </Compile>
+    <Compile Include="command_modules\azure-cli-vm\azure\cli\command_modules\vm\_vm_utils.py">
+      <SubType>Code</SubType>
+    </Compile>
     <Compile Include="command_modules\azure-cli-vm\azure\cli\command_modules\vm\__init__.py" />
     <Compile Include="command_modules\azure-cli-vm\setup.py" />
     <Compile Include="command_modules\azure-cli-webapp\azure\cli\command_modules\webapp\custom.py" />

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

@@ -9,6 +9,7 @@
 from six.moves.urllib.request import urlopen #pylint: disable=import-error
 
 from ._factory import _compute_client_factory, _subscription_client_factory
+from ._vm_utils import read_content_if_is_file
 
 class VMImageFieldAction(argparse.Action): #pylint: disable=too-few-public-methods
     def __call__(self, parser, namespace, values, option_string=None):
@@ -38,13 +39,7 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 class VMSSHFieldAction(argparse.Action): #pylint: disable=too-few-public-methods
     def __call__(self, parser, namespace, values, option_string=None):
-        ssh_value = values
-
-        if os.path.exists(ssh_value):
-            with open(ssh_value, 'r') as f:
-                namespace.ssh_key_value = f.read()
-        else:
-            namespace.ssh_key_value = ssh_value
+        namespace.ssh_key_value = read_content_if_is_file(values)
 
 class VMDNSNameAction(argparse.Action): #pylint: disable=too-few-public-methods
     def __call__(self, parser, namespace, values, option_string=None):

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

@@ -66,4 +66,12 @@
 helps['vm availability-set create'] = """
             type: command
             long-summary: For more info, see https://blogs.technet.microsoft.com/yungchou/2013/05/14/window-azure-fault-domain-and-upgrade-domain-explained-explained-reprised/
+            """
+
+helps['vm extension set'] = """
+            type: command
+            examples:
+                - name: Add a new linux user:
+                  text: 
+                    az vm extension set -n VMAccessForLinux --publisher Microsoft.OSTCExtensions --version 1.4 --vm-name myvm --resource-group yugangw --private-config "{\\"username\\":\\"user1\\", \\"ssh_key\\":\\"ssh_rsa ....\\"}"
             """

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

@@ -0,0 +1,29 @@
+from codecs import open as codecs_open
+import json
+import os
+
+def read_content_if_is_file(string_or_file):
+    content = string_or_file
+    if os.path.exists(string_or_file):
+        with open(string_or_file, 'r') as f:
+            content = f.read()
+    return content
+
+def load_json(string_or_file_path):
+    if os.path.exists(string_or_file_path):
+        return _load_json_from_file(string_or_file_path, 'utf-8') \
+            or _load_json_from_file(string_or_file_path, 'utf-8-sig') \
+            or _load_json_from_file(string_or_file_path, 'utf-16') \
+            or _load_json_from_file(string_or_file_path, 'utf-16le') \
+            or _load_json_from_file(string_or_file_path, 'utf-16be')
+    else:
+        return json.loads(string_or_file_path)
+
+def _load_json_from_file(file_path, encoding):
+    try:
+        with codecs_open(file_path, encoding=encoding) as f:
+            text = f.read()
+        return json.loads(text)
+    except ValueError:
+        pass
+

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

@@ -1,13 +1,17 @@
 # pylint: disable=no-self-use,too-many-arguments
-import os
 import re
+try:
+    from urllib.parse import urlparse
+except ImportError:
+    from urlparse import urlparse # pylint: disable=import-error
 from six.moves.urllib.request import urlopen #pylint: disable=import-error,unused-import
 
 from azure.mgmt.compute.models import DataDisk
 from azure.mgmt.compute.models.compute_management_client_enums import DiskCreateOptionTypes
 from azure.cli.commands import CommandTable, LongRunningOperation
 from azure.cli.commands._command_creation import get_mgmt_service_client, get_data_service_client
 from azure.cli._util import CLIError
+from ._vm_utils import read_content_if_is_file, load_json
 
 from ._actions import (load_images_from_aliases_doc,
                        load_extension_images_thru_services,
@@ -108,7 +112,7 @@ def list_vm_extension_images(self,
         '''vm extension image list
         :param str image_location:Image location
         :param str publisher:Image publisher name
-        :param str type:Image name
+        :param str name:Image name
         :param str version:Image version
         '''
         return load_extension_images_thru_services(publisher,
@@ -259,10 +263,7 @@ def set_linux_user(self,
             protected_settings['password'] = password
 
         if ssh_key_value:
-            if os.path.exists(ssh_key_value):
-                with open(ssh_key_value, 'r') as f:
-                    ssh_key_value = f.read()
-            protected_settings['ssh_key'] = ssh_key_value
+            protected_settings['ssh_key'] = read_content_if_is_file(ssh_key_value)
 
         publisher, extension_name, version, auto_upgrade = _get_access_extension_upgrade_info(
             vm.resources, is_linux=True)
@@ -306,7 +307,6 @@ def delete_linux_user(self,
                                                                   extension_name,
                                                                   ext)
 
-
     def disable_boot_diagnostics(self, resource_group_name, vm_name):
         vm = _vm_get(resource_group_name, vm_name)
         diag_profile = vm.diagnostics_profile
@@ -349,10 +349,6 @@ def enable_boot_diagnostics(self, resource_group_name, vm_name, storage_uri):
     def get_boot_log(self, resource_group_name, vm_name):
         import sys
         import io
-        try:
-            from urllib.parse import urlparse
-        except ImportError:
-            from urlparse import urlparse # pylint: disable=import-error
 
         from azure.mgmt.storage import StorageManagementClient, StorageManagementClientConfiguration
         from azure.storage.blob import BlockBlobService
@@ -406,3 +402,52 @@ def write(self, str_or_bytes):
                     self.out.write(str_or_bytes)
 
         storage_client.get_blob_to_stream(container, blob, StreamWriter(sys.stdout))
+
+    def list_extensions(self, resource_group_name, vm_name):
+        vm = _vm_get(resource_group_name, vm_name)
+        extension_type = 'Microsoft.Compute/virtualMachines/extensions'
+        result = [r for r in vm.resources if r.type == extension_type]
+        return result
+
+    def set_extension(self,
+                      resource_group_name,
+                      vm_name,
+                      vm_extension_name,
+                      publisher,
+                      version,
+                      public_config=None,
+                      private_config=None,
+                      auto_upgrade_minor_version=False):
+        '''create/update extensions for a VM in a resource group'
+        :param vm_name: the name of virtual machine.
+        :param vm_extension_name: the name of the extension
+        :param publisher: the name of extension publisher
+        :param version: the version of extension, must be in the format of "major.minor"
+        :param public_config: public configuration content or a file path
+        :param private_config: private configuration content or a file path
+        :param auto_upgrade_minor_version: auto upgrade to the newer version if available
+        '''
+        vm = _vm_get(resource_group_name, vm_name)
+        client = _compute_client_factory()
+
+        from azure.mgmt.compute.models import VirtualMachineExtension
+
+        protected_settings = load_json(private_config) if not private_config else {}
+        settings = load_json(public_config) if not public_config else None
+
+        #workaround a known issue: the version must only contain "major.minor", even though
+        #"extension image list" gives more detail
+        version = '.'.join(version.split('.')[0:2])
+
+        ext = VirtualMachineExtension(vm.location,#pylint: disable=no-member
+                                      publisher=publisher,
+                                      virtual_machine_extension_type=vm_extension_name,
+                                      protected_settings=protected_settings,
+                                      type_handler_version=version,
+                                      settings=settings,
+                                      auto_upgrade_minor_version=auto_upgrade_minor_version)
+
+        return client.virtual_machine_extensions.create_or_update(resource_group_name,
+                                                                  vm_name,
+                                                                  vm_extension_name,
+                                                                  ext)

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

@@ -84,7 +84,20 @@
         CommandDefinition(VirtualMachineExtensionsOperations.delete, LongRunningOperation(L('Deleting VM extension'), L('VM extension deleted'))),
         CommandDefinition(VirtualMachineExtensionsOperations.get, 'VirtualMachineExtension', command_alias='show'),
     ],
-    command_table, {'vm_extension_name': {'name': '--name -n'}})
+    command_table, patch_aliases(PARAMETER_ALIASES, {
+        'vm_extension_name': {'name': '--name -n'}
+        }))
+
+build_operation(
+    'vm extension', None, ConvenienceVmCommands,
+    [
+        CommandDefinition(ConvenienceVmCommands.set_extension, LongRunningOperation(L('Setting extension'), L('Extension was set')), command_alias='set'),
+        CommandDefinition(ConvenienceVmCommands.list_extensions, '[Extensions]', command_alias='list')
+    ],
+    command_table, patch_aliases(PARAMETER_ALIASES, {
+        'vm_extension_name': {'name': '--name -n'},
+        'auto_upgrade_minor_version': {'action': 'store_true'}
+        }))
 
 build_operation(
     'vm image', 'virtual_machine_images', _compute_client_factory,

src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/command_specs.py

@@ -1,9 +1,12 @@
 # AZURE CLI VM TEST DEFINITIONS
 import json
+import os
 import tempfile
 
 from azure.cli.utils.command_test_script import CommandTestScript, JMESPathComparator
 
+TEST_DIR = os.path.abspath(os.path.join(os.path.abspath(__file__), '..'))
+
 #pylint: disable=method-hidden
 class VMImageListByAliasesScenarioTest(CommandTestScript):
 
@@ -634,6 +637,45 @@ def test_body(self):
         verification = [JMESPathComparator('diagnosticsProfile.bootDiagnostics.enabled', False)]
         self.test('vm show {}'.format(common_part), verification)
 
+
+class VMExtensionInstallTest(CommandTestScript):
+
+    def __init__(self):
+        super(VMExtensionInstallTest, self).__init__(None, self.test_body, None)
+
+    def test_body(self):
+        #pylint: disable=line-too-long
+        publisher = 'Microsoft.OSTCExtensions'
+        extension_name = 'VMAccessForLinux'
+        vm_name = 'yugangw8-1'
+        resource_group = 'yugangw8'
+        public_key = ('ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC8InHIPLAu6lMc0d+5voyXqigZfT5r6fAM1+FQAi+mkPDdk2hNq1BG0Bwfc88G'
+                      'm7BImw8TS+x2bnZmhCbVnHd6BPCDY7a+cHCSqrQMW89Cv6Vl4ueGOeAWHpJTV9CTLVz4IY1x4HBdkLI2lKIHri9+z7NIdvFk7iOk'
+                      'MVGyez5H1xDbF2szURxgc4I2/o5wycSwX+G8DrtsBvWLmFv9YAPx+VkEHQDjR0WWezOjuo1rDn6MQfiKfqAjPuInwNOg5AIxXAOR'
+                      'esrin2PUlArNtdDH1zlvI4RZi36+tJO7mtm3dJiKs4Sj7G6b1CjIU6aaj27MmKy3arIFChYav9yYM3IT')
+        user_name = 'yugangw'
+        config_file_name = 'private_config.json'
+        config = {
+            'username': user_name,
+            'ssh_key': public_key
+            }
+        config_file = os.path.join(TEST_DIR, config_file_name)
+        with open(config_file, 'w') as outfile:
+            json.dump(config, outfile)
+
+        try:
+            set_cmd = ('vm extension set -n {} --publisher {} --version 1.4  --vm-name {} --resource-group {} --private-config {}'
+                       .format(extension_name, publisher, vm_name, resource_group, config_file))
+            self.run(set_cmd)
+            self.test('vm extension show --resource-group {} --vm-name {} --name {}'.format(
+                resource_group, vm_name, extension_name), [
+                    JMESPathComparator('type(@)', 'object'),
+                    JMESPathComparator('name', extension_name),
+                    JMESPathComparator('resourceGroup', resource_group)])
+        finally:
+            os.remove(config_file)
+
+
 ENV_VAR = {}
 
 TEST_DEF = [
@@ -728,5 +770,9 @@ def test_body(self):
     {
         'test_name': 'vm_enable_disable_boot_diagnostic',
         'command': VMBootDiagnostics()
+    },
+    {
+        'test_name': 'vm_extension_install',
+        'command': VMExtensionInstallTest()
     }
 ]