Managed Identity for Arc

Author: rayluo

msal/imds.py

@@ -128,6 +128,18 @@ def _obtain_token(http_client, managed_identity, resource):
             managed_identity,
             resource,
         )
+    if "IDENTITY_ENDPOINT" in os.environ and "IMDS_ENDPOINT" in os.environ:
+        if ManagedIdentity.is_user_assigned(managed_identity):
+            raise ValueError(  # Note: Azure Identity for Python raised exception too
+                "Ignoring managed_identity parameter. "
+                "Azure Arc supports only system-assigned managed identity, "
+                "See also "
+                "https://learn.microsoft.com/en-us/azure/service-fabric/configure-existing-cluster-enable-managed-identity-token-service")
+        return _obtain_token_on_arc(
+            http_client,
+            os.environ["IDENTITY_ENDPOINT"],
+            resource,
+        )
     return _obtain_token_on_azure_vm(http_client, managed_identity, resource)
 
 
@@ -248,6 +260,44 @@ def _obtain_token_on_service_fabric(
         raise
 
 
+def _obtain_token_on_arc(http_client, endpoint, resource):
+    # https://learn.microsoft.com/en-us/azure/azure-arc/servers/managed-identity-authentication
+    logger.debug("Obtaining token via managed identity on Azure Arc")
+    resp = http_client.get(
+        endpoint,
+        params={"api-version": "2020-06-01", "resource": resource},
+        headers={"Metadata": "true"},
+        )
+    www_auth = "www-authenticate"  # Header in lower case
+    challenge = {
+        # Normalized to lowercase, because header names are case-insensitive
+        # https://datatracker.ietf.org/doc/html/rfc7230#section-3.2
+        k.lower(): v for k, v in resp.headers.items() if k.lower() == www_auth
+        }.get(www_auth, "").split("=")  # Output will be ["Basic realm", "content"]
+    if not (  # https://datatracker.ietf.org/doc/html/rfc7617#section-2
+            len(challenge) == 2 and challenge[0].lower() == "basic realm"):
+        raise ValueError("Irrecognizable WWW-Authenticate header: {}".format(resp.headers))
+    with open(challenge[1]) as f:
+        secret = f.read()
+    response = http_client.get(
+        endpoint,
+        params={"api-version": "2020-06-01", "resource": resource},
+        headers={"Metadata": "true", "Authorization": "Basic {}".format(secret)},
+        )
+    payload = json.loads(response.text)
+    if payload.get("access_token") and payload.get("expires_in"):
+        # Example: https://learn.microsoft.com/en-us/azure/azure-arc/servers/media/managed-identity-authentication/bash-token-output-example.png
+        return {
+            "access_token": payload["access_token"],
+            "expires_in": int(payload["expires_in"]),
+            "token_type": payload.get("token_type", "Bearer"),
+            "resource": payload.get("resource"),
+            }
+    return {
+        "error": "invalid_request",
+        "error_description": response.text,
+        }
+
 
 class ManagedIdentityClient(object):
     _instance, _tenant = socket.getfqdn(), "managed_identity"  # Placeholders

tests/http_client.py

@@ -27,7 +27,7 @@ def close(self):  # Not required, but we use it to avoid a warning in unit test
 class MinimalResponse(object):  # Not for production use
     def __init__(self, requests_resp=None, status_code=None, text=None):
         self.status_code = status_code or requests_resp.status_code
-        self.text = text or requests_resp.text
+        self.text = text if text is not None else requests_resp.text
         self._raw_resp = requests_resp
 
     def raise_for_status(self):

tests/test_mi.py

@@ -3,9 +3,9 @@
 import time
 import unittest
 try:
-    from unittest.mock import patch, ANY
+    from unittest.mock import patch, ANY, mock_open
 except:
-    from mock import patch, ANY
+    from mock import patch, ANY, mock_open
 import requests
 
 from tests.http_client import MinimalResponse
@@ -59,7 +59,7 @@ def _test_token_cache(self, app):
 
     def _test_happy_path(self, app, mocked_http):
         result = app.acquire_token(resource="R")
-        mocked_http.assert_called_once()
+        mocked_http.assert_called()
         self.assertEqual({
             "access_token": "AT",
             "expires_in": 1234,
@@ -158,3 +158,23 @@ def test_app_service_error_should_be_normalized(self):
             }, self.app.acquire_token(resource="R"))
             self.assertEqual({}, self.app._token_cache._cache)
 
+
+@patch.dict(os.environ, {
+    "IDENTITY_ENDPOINT": "http://localhost/token",
+    "IMDS_ENDPOINT": "http://localhost",
+})
+class ArcTestCase(ClientTestCase):
+
+    @patch("builtins.open", mock_open(read_data="secret"))
+    def test_happy_path(self):
+        with patch.object(self.app._http_client, "get", side_effect=[
+                MinimalResponse(status_code=401, text="", headers={
+                    "WWW-Authenticate": "Basic realm=/tmp/foo",
+                    }),
+                MinimalResponse(
+                    status_code=200,
+                    text='{"access_token": "AT", "expires_in": "1234", "resource": "R"}',
+                    ),
+                ]) as mocked_method:
+            super(ArcTestCase, self)._test_happy_path(self.app, mocked_method)
+