Add UMI support in MSAL

Author: heyitsaamir

packages/api/src/microsoft/teams/api/auth/__init__.py

@@ -4,7 +4,7 @@
 """
 
 from .caller import CallerIds, CallerType
-from .credentials import ClientCredentials, Credentials, TokenCredentials
+from .credentials import ClientCredentials, Credentials, ManagedIdentityCredentials, TokenCredentials
 from .json_web_token import JsonWebToken, JsonWebTokenPayload
 from .token import TokenProtocol
 
@@ -13,6 +13,7 @@
     "CallerType",
     "ClientCredentials",
     "Credentials",
+    "ManagedIdentityCredentials",
     "TokenCredentials",
     "TokenProtocol",
     "JsonWebToken",

packages/api/src/microsoft/teams/api/auth/credentials.py

@@ -3,7 +3,7 @@
 Licensed under the MIT License.
 """
 
-from typing import Awaitable, Callable, Optional, Union
+from typing import Awaitable, Callable, Literal, Optional, Union
 
 from ..models import CustomBaseModel
 
@@ -43,5 +43,22 @@ class TokenCredentials(CustomBaseModel):
     """
 
 
+class ManagedIdentityCredentials(CustomBaseModel):
+    """Credentials for authentication using Azure Managed Identity."""
+
+    client_id: str
+    """
+    The client ID of the app registration.
+    """
+    managed_identity_type: Literal["system", "user"]
+    """
+    The type of managed identity: 'system' for system-assigned or 'user' for user-assigned.
+    """
+    tenant_id: Optional[str] = None
+    """
+    The tenant ID.
+    """
+
+
 # Union type for credentials
-Credentials = Union[ClientCredentials, TokenCredentials]
+Credentials = Union[ClientCredentials, TokenCredentials, ManagedIdentityCredentials]

packages/apps/src/microsoft/teams/apps/app.py

@@ -21,6 +21,7 @@
     ConversationAccount,
     ConversationReference,
     Credentials,
+    ManagedIdentityCredentials,
     MessageActivityInput,
     TokenCredentials,
 )
@@ -290,13 +291,25 @@ def _init_credentials(self) -> Optional[Credentials]:
         client_secret = self.options.client_secret or os.getenv("CLIENT_SECRET")
         tenant_id = self.options.tenant_id or os.getenv("TENANT_ID")
         token = self.options.token
+        enable_managed_identity = self.options.enable_managed_identity or os.getenv("ENABLE_MANAGED_IDENTITY")
 
         self.log.debug(f"Using CLIENT_ID: {client_id}")
         if not tenant_id:
             self.log.warning("TENANT_ID is not set, assuming multi-tenant app")
         else:
             self.log.debug(f"Using TENANT_ID: {tenant_id} (assuming single-tenant app)")
 
+        if enable_managed_identity and client_id:
+            assert enable_managed_identity in ("system", "user"), (
+                f"enable_managed_identity must be 'system' or 'user', got: {enable_managed_identity}"
+            )
+            self.log.debug(f"Using managed identity: {enable_managed_identity}")
+            return ManagedIdentityCredentials(
+                client_id=client_id,
+                managed_identity_type=enable_managed_identity,
+                tenant_id=tenant_id,
+            )
+
         # - If client_id + client_secret : use ClientCredentials (standard client auth)
         if client_id and client_secret:
             return ClientCredentials(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id)

packages/apps/src/microsoft/teams/apps/options.py

@@ -5,7 +5,7 @@
 
 from dataclasses import dataclass, field
 from logging import Logger
-from typing import Any, Awaitable, Callable, List, Optional, TypedDict, Union, cast
+from typing import Any, Awaitable, Callable, List, Literal, Optional, TypedDict, Union, cast
 
 from microsoft.teams.common import Storage
 from typing_extensions import Unpack
@@ -22,6 +22,8 @@ class AppOptions(TypedDict, total=False):
     tenant_id: Optional[str]
     # Custom token provider function
     token: Optional[Callable[[Union[str, list[str]], Optional[str]], Union[str, Awaitable[str]]]]
+    # Managed identity configuration
+    enable_managed_identity: Optional[Literal["system", "user"]]
 
     # Infrastructure
     logger: Optional[Logger]
@@ -47,6 +49,7 @@ class InternalAppOptions:
     client_secret: Optional[str] = None
     tenant_id: Optional[str] = None
     token: Optional[Callable[[Union[str, list[str]], Optional[str]], Union[str, Awaitable[str]]]] = None
+    enable_managed_identity: Optional[Literal["system", "user"]] = None
     logger: Optional[Logger] = None
     storage: Optional[Storage[str, Any]] = None
 

packages/apps/src/microsoft/teams/apps/token_manager.py

@@ -8,15 +8,21 @@
 from inspect import isawaitable
 from typing import Any, Optional, reveal_type
 
+import requests
 from microsoft.teams.api import (
     ClientCredentials,
     Credentials,
     JsonWebToken,
     TokenProtocol,
 )
-from microsoft.teams.api.auth.credentials import TokenCredentials
+from microsoft.teams.api.auth.credentials import ManagedIdentityCredentials, TokenCredentials
 from microsoft.teams.common import ConsoleLogger
-from msal import ConfidentialClientApplication  # pyright: ignore[reportMissingTypeStubs]
+from msal import (  # pyright: ignore[reportMissingTypeStubs]
+    ConfidentialClientApplication,
+    ManagedIdentityClient,
+    SystemAssignedManagedIdentity,
+    UserAssignedManagedIdentity,
+)
 
 
 class TokenManager:
@@ -36,7 +42,7 @@ def __init__(
         else:
             self._logger = logger.getChild("TokenManager")
 
-        self._msal_clients_by_tenantId: dict[str, ConfidentialClientApplication] = {}
+        self._msal_clients_by_tenantId: dict[str, ConfidentialClientApplication | ManagedIdentityClient] = {}
 
     async def get_bot_token(self) -> Optional[TokenProtocol]:
         """Refresh the bot authentication token."""
@@ -63,9 +69,9 @@ async def _get_token(
             if caller_name:
                 self._logger.debug(f"No credentials provided for {caller_name}")
             return None
-        if isinstance(credentials, ClientCredentials):
+        if isinstance(credentials, (ClientCredentials, ManagedIdentityCredentials)):
             tenant_id_param = tenant_id or credentials.tenant_id or "botframework.com"
-            msal_client = self._get_msal_client_for_tenant(tenant_id_param)
+            msal_client = self._get_msal_client(tenant_id_param)
             token_res: dict[str, Any] | None = await asyncio.to_thread(
                 lambda: msal_client.acquire_token_for_client(scope if isinstance(scope, list) else [scope])
             )
@@ -89,18 +95,34 @@ async def _get_token(
 
             return JsonWebToken(access_token)
 
-    def _get_msal_client_for_tenant(self, tenant_id: str) -> ConfidentialClientApplication:
+    def _get_msal_client(self, tenant_id: str) -> ConfidentialClientApplication | ManagedIdentityClient:
         credentials = self._credentials
-        assert isinstance(credentials, ClientCredentials), (
-            "MSAL clients are only eligible for client credentials,"
-            f"but current credentials is {reveal_type(credentials)}"
-        )
-        cached_client = self._msal_clients_by_tenantId.setdefault(
-            tenant_id,
-            ConfidentialClientApplication(
+
+        # Check if client already exists in cache
+        cached_client = self._msal_clients_by_tenantId.get(tenant_id)
+        if cached_client:
+            return cached_client
+
+        # Create the appropriate client based on credential type
+        if isinstance(credentials, ClientCredentials):
+            client: ConfidentialClientApplication | ManagedIdentityClient = ConfidentialClientApplication(
                 credentials.client_id,
-                client_credential=credentials.client_secret if credentials else None,
+                client_credential=credentials.client_secret,
                 authority=f"https://login.microsoftonline.com/{tenant_id}",
-            ),
-        )
-        return cached_client
+            )
+        elif isinstance(credentials, ManagedIdentityCredentials):
+            # Create the appropriate managed identity based on type
+            if credentials.managed_identity_type == "system":
+                managed_identity = SystemAssignedManagedIdentity()
+            else:  # "user"
+                managed_identity = UserAssignedManagedIdentity(client_id=credentials.client_id)
+
+            client = ManagedIdentityClient(
+                managed_identity,
+                http_client=requests.Session(),
+            )
+        else:
+            raise ValueError(f"Unsupported credential type: {reveal_type(credentials)}")
+
+        self._msal_clients_by_tenantId[tenant_id] = client
+        return client

stubs/msal/__init__.pyi

@@ -11,3 +11,29 @@ class ConfidentialClientApplication:
     def acquire_token_for_client(
         self, scopes: list[str] | str, claims_challenge: Optional[str] = None, **kwargs: Any
     ) -> dict[str, Any]: ...
+
+class SystemAssignedManagedIdentity:
+    """MSAL System Assigned Managed Identity"""
+
+    def __init__(self) -> None: ...
+
+class UserAssignedManagedIdentity:
+    """MSAL User Assigned Managed Identity"""
+
+    def __init__(self, *, client_id: str) -> None: ...
+
+class ManagedIdentityClient:
+    """MSAL Managed Identity Client"""
+
+    def __init__(
+        self,
+        managed_identity: SystemAssignedManagedIdentity | UserAssignedManagedIdentity | dict[str, Any],
+        *,
+        http_client: Any,
+        token_cache: Optional[Any] = None,
+        http_cache: Optional[Any] = None,
+        client_capabilities: Optional[list[str]] = None,
+    ) -> None: ...
+    def acquire_token_for_client(
+        self, scopes: list[str] | str, claims_challenge: Optional[str] = None, **kwargs: Any
+    ) -> dict[str, Any]: ...