[Profile] az login: Support --claims-challenge in auth code flow (#31778)

Author: jiasli

src/azure-cli-core/azure/cli/core/_profile.py

@@ -152,7 +152,7 @@ def login(self,
               allow_no_subscriptions=False,
               use_cert_sn_issuer=None,
               show_progress=False,
-              **kwargs):
+              claims_challenge=None):
         """
         For service principal, `password` is a dict returned by ServicePrincipalAuth.build_credential
         """
@@ -172,12 +172,12 @@ def login(self,
                 use_device_code = True
 
             if use_device_code:
-                user_identity = identity.login_with_device_code(scopes=scopes, **kwargs)
+                user_identity = identity.login_with_device_code(scopes=scopes)
             else:
-                user_identity = identity.login_with_auth_code(scopes=scopes, **kwargs)
+                user_identity = identity.login_with_auth_code(scopes=scopes, claims_challenge=claims_challenge)
         else:
             if not is_service_principal:
-                user_identity = identity.login_with_username_password(username, password, scopes=scopes, **kwargs)
+                user_identity = identity.login_with_username_password(username, password, scopes=scopes)
             else:
                 identity.login_with_service_principal(username, password, scopes=scopes)
 

src/azure-cli-core/azure/cli/core/auth/identity.py

@@ -145,7 +145,7 @@ def _service_principal_store(self):
             Identity._service_principal_store_instance = ServicePrincipalStore(store)
         return Identity._service_principal_store_instance
 
-    def login_with_auth_code(self, scopes, **kwargs):
+    def login_with_auth_code(self, scopes, claims_challenge=None):
         # Emit a warning to inform that a browser is opened.
         # Only show the path part of the URL and hide the query string.
 
@@ -168,21 +168,21 @@ def _prompt_launching_ui(ui=None, **_):
             success_template=success_template, error_template=error_template,
             parent_window_handle=self._msal_app.CONSOLE_WINDOW_HANDLE, on_before_launching_ui=_prompt_launching_ui,
             enable_msa_passthrough=True,
-            **kwargs)
+            claims_challenge=claims_challenge)
         return check_result(result)
 
-    def login_with_device_code(self, scopes, **kwargs):
-        flow = self._msal_app.initiate_device_flow(scopes, **kwargs)
+    def login_with_device_code(self, scopes):
+        flow = self._msal_app.initiate_device_flow(scopes)
         if "user_code" not in flow:
             raise ValueError(
                 "Fail to create device flow. Err: %s" % json.dumps(flow, indent=4))
         from azure.cli.core.style import print_styled_text, Style
         print_styled_text((Style.WARNING, flow["message"]), file=sys.stderr)
-        result = self._msal_app.acquire_token_by_device_flow(flow, **kwargs)  # By default it will block
+        result = self._msal_app.acquire_token_by_device_flow(flow)  # By default it will block
         return check_result(result)
 
-    def login_with_username_password(self, username, password, scopes, **kwargs):
-        result = self._msal_app.acquire_token_by_username_password(username, password, scopes, **kwargs)
+    def login_with_username_password(self, username, password, scopes):
+        result = self._msal_app.acquire_token_by_username_password(username, password, scopes)
         return check_result(result)
 
     def login_with_service_principal(self, client_id, credential, scopes):

src/azure-cli-core/azure/cli/core/util.py

@@ -603,13 +603,24 @@ def shell_safe_json_parse(json_or_dict_string, preserve_order=False, strict=True
 
 def b64encode(s):
     """
-    Encodes a string to base64 on 2.x and 3.x
+    Encodes a string to a base64 string.
     :param str s: latin_1 encoded string
     :return: base64 encoded string
     :rtype: str
     """
     encoded = base64.b64encode(s.encode("latin-1"))
-    return encoded if encoded is str else encoded.decode('latin-1')
+    return encoded.decode('latin-1')
+
+
+def b64decode(s):
+    """
+    Decodes a base64 string to a string.
+    :param str s: latin_1 encoded base64 string
+    :return: decoded string
+    :rtype: str
+    """
+    encoded = base64.b64decode(s.encode("latin-1"))
+    return encoded.decode('latin-1')
 
 
 def b64_to_hex(s):

src/azure-cli/azure/cli/command_modules/profile/__init__.py

@@ -55,6 +55,9 @@ def load_arguments(self, command):
             c.argument('allow_no_subscriptions', action='store_true',
                        help="Support accessing tenants without subscriptions. It's useful to run "
                             "tenant-level commands, such as 'az ad'.")
+            c.argument('claims_challenge',
+                       help="Base64-encoded claims challenge requested by a resource API in the "
+                            "WWW-Authenticate header.")
             c.ignore('_subscription')  # hide the global subscription parameter
 
             # Device code flow

src/azure-cli/azure/cli/command_modules/profile/custom.py

@@ -122,6 +122,7 @@ def account_clear(cmd):
 
 # pylint: disable=too-many-branches, too-many-locals
 def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_subscriptions=False,
+          claims_challenge=None,
           # Device code flow
           use_device_code=False,
           # Service principal
@@ -148,6 +149,10 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_
         else:
             logger.warning(USERNAME_PASSWORD_DEPRECATION_WARNING_OTHER_CLOUD)
 
+    if claims_challenge:
+        from azure.cli.core.util import b64decode
+        claims_challenge = b64decode(claims_challenge)
+
     interactive = False
 
     profile = Profile(cli_ctx=cmd.cli_ctx)
@@ -194,7 +199,9 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_
         use_device_code=use_device_code,
         allow_no_subscriptions=allow_no_subscriptions,
         use_cert_sn_issuer=use_cert_sn_issuer,
-        show_progress=select_subscription)
+        show_progress=select_subscription,
+        claims_challenge=claims_challenge
+    )
 
     # Launch interactive account selection. No JSON output.
     if select_subscription: