|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import logging |
| 4 | +import jwt |
4 | 5 | from typing import Optional |
5 | 6 | from urllib.parse import urlparse, ParseResult as URI |
6 | 7 | from msal import ( |
|
23 | 24 | logger = logging.getLogger(__name__) |
24 | 25 |
|
25 | 26 |
|
| 27 | +# this is deferred because jwt.decode is expensive and we don't want to do it unless we |
| 28 | +# have logging.DEBUG enabled |
| 29 | +class _DeferredLogOfBlueprintId: |
| 30 | + def __init__(self, jwt_token: str): |
| 31 | + self.jwt_token = jwt_token |
| 32 | + |
| 33 | + def __str__(self): |
| 34 | + payload = jwt.decode(self.jwt_token, options={"verify_signature": False}) |
| 35 | + agentic_blueprint_id = payload.get("xms_par_app_azp") |
| 36 | + return f"Agentic blueprint id: {agentic_blueprint_id}" |
| 37 | + |
| 38 | + |
26 | 39 | class MsalAuth(AccessTokenProviderBase): |
27 | 40 |
|
28 | 41 | _client_credential_cache = None |
@@ -56,11 +69,16 @@ async def get_access_token( |
56 | 69 | auth_result_payload = msal_auth_client.acquire_token_for_client( |
57 | 70 | scopes=local_scopes |
58 | 71 | ) |
| 72 | + else: |
| 73 | + auth_result_payload = None |
59 | 74 |
|
60 | | - # TODO: Handling token error / acquisition failed |
61 | | - return auth_result_payload["access_token"] |
| 75 | + res = auth_result_payload.get("access_token") if auth_result_payload else None |
| 76 | + if not res: |
| 77 | + logger.error("Failed to acquire token for resource %s", auth_result_payload) |
| 78 | + raise ValueError(f"Failed to acquire token. {str(auth_result_payload)}") |
| 79 | + return res |
62 | 80 |
|
63 | | - async def aquire_token_on_behalf_of( |
| 81 | + async def acquire_token_on_behalf_of( |
64 | 82 | self, scopes: list[str], user_assertion: str |
65 | 83 | ) -> str: |
66 | 84 | """ |
@@ -186,3 +204,189 @@ def _resolve_scopes_list(self, instance_url: URI, scopes=None) -> list[str]: |
186 | 204 | temp_list.append(scope_placeholder) |
187 | 205 | logger.debug(f"Resolved scopes: {temp_list}") |
188 | 206 | return temp_list |
| 207 | + |
| 208 | + # the call to MSAL is blocking, but in the future we want to create an asyncio task |
| 209 | + # to avoid this |
| 210 | + async def get_agentic_application_token( |
| 211 | + self, agent_app_instance_id: str |
| 212 | + ) -> Optional[str]: |
| 213 | + """Gets the agentic application token for the given agent application instance ID. |
| 214 | +
|
| 215 | + :param agent_app_instance_id: The agent application instance ID. |
| 216 | + :type agent_app_instance_id: str |
| 217 | + :return: The agentic application token, or None if not found. |
| 218 | + :rtype: Optional[str] |
| 219 | + """ |
| 220 | + |
| 221 | + if not agent_app_instance_id: |
| 222 | + raise ValueError("Agent application instance Id must be provided.") |
| 223 | + |
| 224 | + logger.info( |
| 225 | + "Attempting to get agentic application token from agent_app_instance_id %s", |
| 226 | + agent_app_instance_id, |
| 227 | + ) |
| 228 | + msal_auth_client = self._create_client_application() |
| 229 | + |
| 230 | + if isinstance(msal_auth_client, ConfidentialClientApplication): |
| 231 | + |
| 232 | + # https://github.dev/AzureAD/microsoft-authentication-library-for-dotnet |
| 233 | + auth_result_payload = msal_auth_client.acquire_token_for_client( |
| 234 | + ["api://AzureAdTokenExchange/.default"], |
| 235 | + data={"fmi_path": agent_app_instance_id}, |
| 236 | + ) |
| 237 | + |
| 238 | + if auth_result_payload: |
| 239 | + return auth_result_payload.get("access_token") |
| 240 | + |
| 241 | + return None |
| 242 | + |
| 243 | + async def get_agentic_instance_token( |
| 244 | + self, agent_app_instance_id: str |
| 245 | + ) -> tuple[str, str]: |
| 246 | + """Gets the agentic instance token for the given agent application instance ID. |
| 247 | +
|
| 248 | + :param agent_app_instance_id: The agent application instance ID. |
| 249 | + :type agent_app_instance_id: str |
| 250 | + :return: A tuple containing the agentic instance token and the agent application token. |
| 251 | + :rtype: tuple[str, str] |
| 252 | + """ |
| 253 | + |
| 254 | + if not agent_app_instance_id: |
| 255 | + raise ValueError("Agent application instance Id must be provided.") |
| 256 | + |
| 257 | + logger.info( |
| 258 | + "Attempting to get agentic instance token from agent_app_instance_id %s", |
| 259 | + agent_app_instance_id, |
| 260 | + ) |
| 261 | + agent_token_result = await self.get_agentic_application_token( |
| 262 | + agent_app_instance_id |
| 263 | + ) |
| 264 | + |
| 265 | + if not agent_token_result: |
| 266 | + logger.error( |
| 267 | + "Failed to acquire agentic instance token or agent token for agent_app_instance_id %s", |
| 268 | + agent_app_instance_id, |
| 269 | + ) |
| 270 | + raise Exception( |
| 271 | + f"Failed to acquire agentic instance token or agent token for agent_app_instance_id {agent_app_instance_id}" |
| 272 | + ) |
| 273 | + |
| 274 | + authority = ( |
| 275 | + f"https://login.microsoftonline.com/{self._msal_configuration.TENANT_ID}" |
| 276 | + ) |
| 277 | + |
| 278 | + instance_app = ConfidentialClientApplication( |
| 279 | + client_id=agent_app_instance_id, |
| 280 | + authority=authority, |
| 281 | + client_credential={"client_assertion": agent_token_result}, |
| 282 | + ) |
| 283 | + |
| 284 | + agentic_instance_token = instance_app.acquire_token_for_client( |
| 285 | + ["api://AzureAdTokenExchange/.default"] |
| 286 | + ) |
| 287 | + |
| 288 | + if not agentic_instance_token: |
| 289 | + logger.error( |
| 290 | + "Failed to acquire agentic instance token or agent token for agent_app_instance_id %s", |
| 291 | + agent_app_instance_id, |
| 292 | + ) |
| 293 | + raise Exception( |
| 294 | + f"Failed to acquire agentic instance token or agent token for agent_app_instance_id {agent_app_instance_id}" |
| 295 | + ) |
| 296 | + |
| 297 | + # future scenario where we don't know the blueprint id upfront |
| 298 | + |
| 299 | + token = agentic_instance_token.get("access_token") |
| 300 | + if not token: |
| 301 | + logger.error( |
| 302 | + "Failed to acquire agentic instance token, %s", agentic_instance_token |
| 303 | + ) |
| 304 | + raise ValueError(f"Failed to acquire token. {str(agentic_instance_token)}") |
| 305 | + |
| 306 | + logger.debug(_DeferredLogOfBlueprintId(token)) |
| 307 | + |
| 308 | + return agentic_instance_token["access_token"], agent_token_result |
| 309 | + |
| 310 | + async def get_agentic_user_token( |
| 311 | + self, agent_app_instance_id: str, upn: str, scopes: list[str] |
| 312 | + ) -> Optional[str]: |
| 313 | + """Gets the agentic user token for the given agent application instance ID and user principal name and the scopes. |
| 314 | +
|
| 315 | + :param agent_app_instance_id: The agent application instance ID. |
| 316 | + :type agent_app_instance_id: str |
| 317 | + :param upn: The user principal name. |
| 318 | + :type upn: str |
| 319 | + :param scopes: The scopes to request for the token. |
| 320 | + :type scopes: list[str] |
| 321 | + :return: The agentic user token, or None if not found. |
| 322 | + :rtype: Optional[str] |
| 323 | + """ |
| 324 | + if not agent_app_instance_id or not upn: |
| 325 | + raise ValueError( |
| 326 | + "Agent application instance Id and user principal name must be provided." |
| 327 | + ) |
| 328 | + |
| 329 | + logger.info( |
| 330 | + "Attempting to get agentic user token from agent_app_instance_id %s and upn %s", |
| 331 | + agent_app_instance_id, |
| 332 | + upn, |
| 333 | + ) |
| 334 | + instance_token, agent_token = await self.get_agentic_instance_token( |
| 335 | + agent_app_instance_id |
| 336 | + ) |
| 337 | + |
| 338 | + if not instance_token or not agent_token: |
| 339 | + logger.error( |
| 340 | + "Failed to acquire instance token or agent token for agent_app_instance_id %s and upn %s", |
| 341 | + agent_app_instance_id, |
| 342 | + upn, |
| 343 | + ) |
| 344 | + raise Exception( |
| 345 | + f"Failed to acquire instance token or agent token for agent_app_instance_id {agent_app_instance_id} and upn {upn}" |
| 346 | + ) |
| 347 | + |
| 348 | + authority = ( |
| 349 | + f"https://login.microsoftonline.com/{self._msal_configuration.TENANT_ID}" |
| 350 | + ) |
| 351 | + |
| 352 | + instance_app = ConfidentialClientApplication( |
| 353 | + client_id=agent_app_instance_id, |
| 354 | + authority=authority, |
| 355 | + client_credential={"client_assertion": agent_token}, |
| 356 | + ) |
| 357 | + |
| 358 | + logger.info( |
| 359 | + "Acquiring agentic user token for agent_app_instance_id %s and upn %s", |
| 360 | + agent_app_instance_id, |
| 361 | + upn, |
| 362 | + ) |
| 363 | + auth_result_payload = instance_app.acquire_token_for_client( |
| 364 | + scopes, |
| 365 | + data={ |
| 366 | + "username": upn, |
| 367 | + "user_federated_identity_credential": instance_token, |
| 368 | + "grant_type": "user_fic", |
| 369 | + }, |
| 370 | + ) |
| 371 | + |
| 372 | + if not auth_result_payload: |
| 373 | + logger.error( |
| 374 | + "Failed to acquire agentic user token for agent_app_instance_id %s and upn %s, %s", |
| 375 | + agent_app_instance_id, |
| 376 | + upn, |
| 377 | + auth_result_payload, |
| 378 | + ) |
| 379 | + return None |
| 380 | + |
| 381 | + access_token = auth_result_payload.get("access_token") |
| 382 | + if not access_token: |
| 383 | + logger.error( |
| 384 | + "Failed to acquire agentic user token for agent_app_instance_id %s and upn %s, %s", |
| 385 | + agent_app_instance_id, |
| 386 | + upn, |
| 387 | + auth_result_payload, |
| 388 | + ) |
| 389 | + return None |
| 390 | + |
| 391 | + logger.info("Acquired agentic user token response.") |
| 392 | + return access_token |
0 commit comments