PoC: Managed Identity for App Service and Azure Automation

Author: rayluo

msal/imds.py

@@ -4,6 +4,8 @@
 # This code is licensed under the MIT License.
 import json
 import logging
+import os
+import time
 try:  # Python 2
     from urlparse import urlparse
 except:  # Python 3
@@ -17,8 +19,18 @@ def _scope_to_resource(scope):  # This is an experimental reasonable-effort appr
         return "{}://{}".format(u.scheme, u.netloc)
     return scope  # There is no much else we can do here
 
+
 def _obtain_token(http_client, resource, client_id=None):
+    if "IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ:
+        return _obtain_token_on_app_service(
+            http_client, os.environ["IDENTITY_ENDPOINT"], os.environ["IDENTITY_HEADER"],
+            resource, client_id=client_id)
+    return _obtain_token_on_azure_vm(http_client, resource, client_id=client_id)
+
+
+def _obtain_token_on_azure_vm(http_client, resource, client_id=None):
     # Based on https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http
+    logger.debug("Obtaining token via managed identity on Azure VM")
     params = {
         "api-version": "2018-02-01",
         "resource": resource,
@@ -44,3 +56,42 @@ def _obtain_token(http_client, resource, client_id=None):
         logger.debug("IMDS emits unexpected payload: %s", resp.text)
         raise
 
+def _obtain_token_on_app_service(http_client, endpoint, identity_header, resource, client_id=None):
+    # Prerequisite: Create your app service https://docs.microsoft.com/en-us/azure/app-service/quickstart-python
+    # Assign it a managed identity https://docs.microsoft.com/en-us/azure/app-service/overview-managed-identity?tabs=portal%2Chttp
+    # SSH into your container for testing https://docs.microsoft.com/en-us/azure/app-service/configure-linux-open-ssh-session
+    logger.debug("Obtaining token via managed identity on Azure App Service")
+    params = {
+        "api-version": "2019-08-01",
+        "resource": resource,
+        }
+    if client_id:
+        params["client_id"] = client_id
+    resp = http_client.get(
+        endpoint,
+        params=params,
+        headers={
+            "X-IDENTITY-HEADER": identity_header,
+            "Metadata": "true",  # Unnecessary yet harmless for App Service,
+            # It will be needed by Azure Automation 
+            # https://docs.microsoft.com/en-us/azure/automation/enable-managed-identity-for-automation#get-access-token-for-system-assigned-managed-identity-using-http-get
+            },
+        )
+    try:
+        payload = json.loads(resp.text)
+        if payload.get("access_token") and payload.get("expires_on"):
+            return {  # Normalizing the payload into OAuth2 format
+                "access_token": payload["access_token"],
+                "expires_in": int(payload["expires_on"]) - int(time.time()),
+                "resource": payload.get("resource"),
+                "token_type": payload.get("token_type", "Bearer"),
+                }
+        return {
+            "error": "invalid_scope",  # Empirically, wrong resource ends up with a vague statusCode=500
+            "error_description": "{}, {}".format(
+                payload.get("statusCode"), payload.get("message")),
+            }
+    except ValueError:
+        logger.debug("IMDS emits unexpected payload: %s", resp.text)
+        raise
+

tests/msaltest.py

@@ -161,7 +161,7 @@ def main():
         {"client_id": AZURE_CLI, "name": "Azure CLI (Correctly configured for MSA-PT)"},
         {"client_id": VISUAL_STUDIO, "name": "Visual Studio (Correctly configured for MSA-PT)"},
         {"client_id": "95de633a-083e-42f5-b444-a4295d8e9314", "name": "Whiteboard Services (Non MSA-PT app. Accepts AAD & MSA accounts.)"},
-        {"client_id": None, "client_secret": None, "name": "System-assigned Managed Identity (Only works when running inside a supported environment, such as Azure VM)"},
+        {"client_id": None, "client_secret": None, "name": "System-assigned Managed Identity (Only works when running inside a supported environment, such as Azure VM, Azure App Service, Azure Automation)"},
         ],
         option_renderer=lambda a: a["name"],
         header="Impersonate this app (or you can type in the client_id of your own app)",