|
5 | 5 | import os |
6 | 6 | import uuid |
7 | 7 | from binascii import hexlify |
| 8 | +from contextlib import contextmanager |
8 | 9 | from typing import Any, NamedTuple, Optional |
9 | 10 |
|
10 | 11 | import jwt |
@@ -50,11 +51,30 @@ def __init__(self): |
50 | 51 | # Reset the auth info |
51 | 52 | self.app: msal.ClientApplication = None |
52 | 53 | self._auth_info = {} |
| 54 | + self._use_device_code = False |
53 | 55 |
|
54 | 56 | # Load the auth info and environment variables |
55 | 57 | self._load_auth() |
56 | 58 | self._load_env() |
57 | 59 |
|
| 60 | + @contextmanager |
| 61 | + def device_code_flow(self): |
| 62 | + """Context manager to enable device code authentication flow. |
| 63 | +
|
| 64 | + Temporarily enables device code flow for token acquisition within |
| 65 | + the context block. Automatically resets the flag on exit, even if |
| 66 | + an exception occurs. |
| 67 | +
|
| 68 | + Usage: |
| 69 | + with FabAuth().device_code_flow(): |
| 70 | + FabAuth().get_access_token(scope=...) |
| 71 | + """ |
| 72 | + self._use_device_code = True |
| 73 | + try: |
| 74 | + yield |
| 75 | + finally: |
| 76 | + self._use_device_code = False |
| 77 | + |
58 | 78 | def _save_auth(self): |
59 | 79 | from fabric_cli.utils.fab_secure_io import write_restricted_file |
60 | 80 |
|
@@ -270,6 +290,51 @@ def _get_app(self) -> msal.ClientApplication: |
270 | 290 | ) |
271 | 291 | return self.app |
272 | 292 |
|
| 293 | + def _acquire_token_by_device_code(self, scope: list[str]) -> Optional[dict]: |
| 294 | + """ |
| 295 | + Acquire a token using the device code flow. |
| 296 | +
|
| 297 | + This initiates a device code flow, prints the user code and verification |
| 298 | + URL to the console, and waits for the user to complete authentication |
| 299 | + on another device. If an account already exists in the cache (from a |
| 300 | + prior device code flow in the same session), attempts a silent refresh |
| 301 | + first to avoid prompting the user multiple times for different scopes. |
| 302 | +
|
| 303 | + Args: |
| 304 | + scope: The scopes for which to request the token |
| 305 | +
|
| 306 | + Returns: |
| 307 | + Full MSAL result dictionary containing access_token, expires_on, etc., |
| 308 | + or None if the device flow returns no result. Raises: |
| 309 | + FabricCLIError: When device code flow initiation or token acquisition fails |
| 310 | + """ |
| 311 | + # If we already have an account from a prior device code auth in this |
| 312 | + # session, try silent acquisition with force_refresh to use the refresh |
| 313 | + # token for a different resource scope. |
| 314 | + accounts = self._get_app().get_accounts() |
| 315 | + if accounts: |
| 316 | + token = self._get_app().acquire_token_silent( |
| 317 | + scopes=scope, account=accounts[0], force_refresh=True |
| 318 | + ) |
| 319 | + if token and token.get("access_token"): |
| 320 | + return token |
| 321 | + |
| 322 | + flow = self._get_app().initiate_device_flow(scopes=scope) |
| 323 | + if "user_code" not in flow: |
| 324 | + raise FabricCLIError( |
| 325 | + ErrorMessages.Auth.device_code_flow_failed( |
| 326 | + flow.get("error_description", "Unknown error") |
| 327 | + ), |
| 328 | + con.ERROR_AUTHENTICATION_FAILED, |
| 329 | + ) |
| 330 | + |
| 331 | + # Display the device code message to the user |
| 332 | + utils_ui.print_grey(flow["message"]) |
| 333 | + |
| 334 | + # Wait for the user to authenticate |
| 335 | + token = self._get_app().acquire_token_by_device_flow(flow) |
| 336 | + return token |
| 337 | + |
273 | 338 | def _get_access_token_from_env_vars_if_exist(self, scope): |
274 | 339 | if "FAB_TOKEN" in os.environ and "FAB_TOKEN_ONELAKE" in os.environ: |
275 | 340 | match scope: |
@@ -494,14 +559,19 @@ def acquire_token(self, scope: list[str], interactive_renew=True) -> dict: |
494 | 559 | scopes=scope, account=account |
495 | 560 | ) |
496 | 561 |
|
497 | | - if token is None and interactive_renew: |
498 | | - token = self._get_app().acquire_token_interactive( |
499 | | - scopes=scope, |
500 | | - prompt="select_account", |
501 | | - parent_window_handle=msal.PublicClientApplication.CONSOLE_WINDOW_HANDLE, |
502 | | - ) |
| 562 | + if (token is None or token.get("error")) and interactive_renew: |
| 563 | + if self._use_device_code: |
| 564 | + token = self._acquire_token_by_device_code(scope) |
| 565 | + else: |
| 566 | + token = self._get_app().acquire_token_interactive( |
| 567 | + scopes=scope, |
| 568 | + prompt="select_account", |
| 569 | + parent_window_handle=msal.PublicClientApplication.CONSOLE_WINDOW_HANDLE, |
| 570 | + ) |
503 | 571 | if token is not None and "id_token_claims" in token: |
504 | | - self.set_tenant(token.get("id_token_claims")["tid"]) |
| 572 | + id_token_claims = token.get("id_token_claims") |
| 573 | + if id_token_claims: |
| 574 | + self.set_tenant(id_token_claims["tid"]) |
505 | 575 |
|
506 | 576 | if token and token.get("error"): |
507 | 577 | fab_logger.log_debug( |
|
0 commit comments