From 8c819dd9cb5e97e62aba3623a7dc416017f862b4 Mon Sep 17 00:00:00 2001
From: Yuwei Zhou <yuwzho@microsoft.com>
Date: Thu, 22 Mar 2018 04:01:44 +0800
Subject: [PATCH] Support MSI for ansible on Azure resources (#36634)

* msi

* add env and param

* add msi in default

* add azure_rm

* add document

* subscription param

* if not enabled msi

* remove the msi in default mode since the infinite loop will block if not enabled msi

* lint

* lint

* Update azure_rm_common.py

* fix

* catch exc and make error message more friendly

* lint

* Minor docs changes to the msi source option
---
 lib/ansible/module_utils/azure_rm_common.py   | 25 ++++++++++++++++++-
 .../utils/module_docs_fragments/azure.py      |  5 ++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/lib/ansible/module_utils/azure_rm_common.py b/lib/ansible/module_utils/azure_rm_common.py
index 0c0d726e31..aef89d2fc3 100644
--- a/lib/ansible/module_utils/azure_rm_common.py
+++ b/lib/ansible/module_utils/azure_rm_common.py
@@ -23,7 +23,7 @@ except ImportError:
 AZURE_COMMON_ARGS = dict(
     auth_source=dict(
         type='str',
-        choices=['auto', 'cli', 'env', 'credential_file']
+        choices=['auto', 'cli', 'env', 'credential_file', 'msi']
     ),
     profile=dict(type='str'),
     subscription_id=dict(type='str', no_log=True),
@@ -127,6 +127,7 @@ except ImportError as exc:
 try:
     from enum import Enum
     from msrestazure.azure_exceptions import CloudError
+    from msrestazure.azure_active_directory import MSIAuthentication
     from msrestazure.tools import resource_id, is_valid_resource_id
     from msrestazure import azure_cloud
     from azure.common.credentials import ServicePrincipalCredentials, UserPassCredentials
@@ -138,6 +139,7 @@ try:
     from azure.mgmt.web.version import VERSION as web_client_version
     from azure.mgmt.network import NetworkManagementClient
     from azure.mgmt.resource.resources import ResourceManagementClient
+    from azure.mgmt.resource.subscriptions import SubscriptionClient
     from azure.mgmt.storage import StorageManagementClient
     from azure.mgmt.compute import ComputeManagementClient
     from azure.mgmt.dns import DnsManagementClient
@@ -489,6 +491,23 @@ class AzureRMModuleBase(object):
 
         return None
 
+    def _get_msi_credentials(self, subscription_id_param=None):
+        credentials = MSIAuthentication()
+        subscription_id = subscription_id_param or os.environ.get(AZURE_CREDENTIAL_ENV_MAPPING['subscription_id'], None)
+        if not subscription_id:
+            try:
+                # use the first subscription of the MSI
+                subscription_client = SubscriptionClient(credentials)
+                subscription = next(subscription_client.subscriptions.list())
+                subscription_id = str(subscription.subscription_id)
+            except Exception as exc:
+                self.fail("Failed to get MSI token: {0}. "
+                          "Please check whether your machine enabled MSI or grant access to any subscription.".format(str(exc)))
+        return {
+            'credentials': credentials,
+            'subscription_id': subscription_id
+        }
+
     def _get_azure_cli_credentials(self):
         credentials, subscription_id = get_azure_cli_credentials()
         cloud_environment = get_cli_active_cloud()
@@ -526,6 +545,10 @@ class AzureRMModuleBase(object):
         if not auth_source:
             auth_source = os.environ.get('ANSIBLE_AZURE_AUTH_SOURCE', 'auto')
 
+        if auth_source == 'msi':
+            self.log('Retrieving credenitals from MSI')
+            return self._get_msi_credentials(arg_credentials['subscription_id'])
+
         if auth_source == 'cli':
             if not HAS_AZURE_CLI_CORE:
                 self.fail("Azure auth_source is `cli`, but azure-cli package is not available. Try `pip install azure-cli --upgrade`")
diff --git a/lib/ansible/utils/module_docs_fragments/azure.py b/lib/ansible/utils/module_docs_fragments/azure.py
index ad8a882525..c826d533aa 100644
--- a/lib/ansible/utils/module_docs_fragments/azure.py
+++ b/lib/ansible/utils/module_docs_fragments/azure.py
@@ -68,11 +68,16 @@ options:
               C(~/.azure/credentials).
             - When set to C(cli), the credentials will be sources from the default Azure CLI profile.
             - Can also be set via the C(ANSIBLE_AZURE_AUTH_SOURCE) environment variable.
+            - When set to C(msi), the host machine must be an azure resource with an enabled MSI extension. C(subscription_id) or the
+              environment variable C(AZURE_SUBSCRIPTION_ID) can be used to identify the subscription ID if the resource is granted
+              access to more than one subscription, otherwise the first subscription is chosen.
+            - The C(msi) was added in Ansible 2.6.
         choices:
         - auto
         - cli
         - credential_file
         - env
+        - msi
         default: auto
         version_added: 2.5
     api_profile: