Federated flow fix (#240)

Author: abhidnya13

adal/token_request.py

@@ -190,9 +190,9 @@ def _perform_wstrust_assertion_oauth_exchange(self, wstrust_response):
 
         return self._oauth_get_token(oauth_parameters)
 
-    def _perform_wstrust_exchange(self, wstrust_endpoint, wstrust_endpoint_version, username, password):
+    def _perform_wstrust_exchange(self, wstrust_endpoint, wstrust_endpoint_version, cloud_audience_urn, username, password):
 
-        wstrust = self._create_wstrust_request(wstrust_endpoint, "urn:federation:MicrosoftOnline",
+        wstrust = self._create_wstrust_request(wstrust_endpoint, cloud_audience_urn,
                                                wstrust_endpoint_version)
         result = wstrust.acquire_token(username, password)
 
@@ -204,15 +204,16 @@ def _perform_wstrust_exchange(self, wstrust_endpoint, wstrust_endpoint_version,
 
         return result
 
-    def _perform_username_password_for_access_token_exchange(self, wstrust_endpoint, wstrust_endpoint_version,
+    def _perform_username_password_for_access_token_exchange(self, wstrust_endpoint, wstrust_endpoint_version, cloud_audience_urn,
                                                              username, password):
-        wstrust_response = self._perform_wstrust_exchange(wstrust_endpoint, wstrust_endpoint_version,
+        wstrust_response = self._perform_wstrust_exchange(wstrust_endpoint, wstrust_endpoint_version, cloud_audience_urn,
                                                           username, password)
         return self._perform_wstrust_assertion_oauth_exchange(wstrust_response)
 
     def _get_token_username_password_federated(self, username, password):
         self._log.debug("Acquiring token with username password for federated user")
 
+        cloud_audience_urn = self._user_realm.cloud_audience_urn
         if not self._user_realm.federation_metadata_url:
             self._log.warn("Unable to retrieve federationMetadataUrl from AAD. "
                            "Attempting fallback to AAD supplied endpoint.")
@@ -228,7 +229,7 @@ def _get_token_username_password_federated(self, username, password):
 
             return self._perform_username_password_for_access_token_exchange(
                 self._user_realm.federation_active_auth_url,
-                wstrust_version, username, password)
+                wstrust_version, cloud_audience_urn, username, password)
         else:
             mex_endpoint = self._user_realm.federation_metadata_url
             self._log.debug(
@@ -253,6 +254,7 @@ def _get_token_username_password_federated(self, username, password):
                     raise AdalError('AAD did not return a WSTrust endpoint. Unable to proceed.')
 
             return self._perform_username_password_for_access_token_exchange(wstrust_endpoint, wstrust_version,
+                                                                             cloud_audience_urn,
                                                                              username, password)
     @staticmethod
     def _parse_wstrust_version_from_federation_active_authurl(federation_active_authurl):

adal/user_realm.py

@@ -57,6 +57,7 @@ def __init__(self, call_context, user_principle, authority_url):
         self.account_type = None
         self.federation_metadata_url = None
         self.federation_active_auth_url = None
+        self.cloud_audience_urn = None
         self._user_principle = user_principle
         self._authority_url = authority_url
 
@@ -131,6 +132,7 @@ def _parse_discovery_response(self, body):
             self.federation_protocol = protocol
             self.federation_metadata_url = response['federation_metadata_url']
             self.federation_active_auth_url = response['federation_active_auth_url']
+            self.cloud_audience_urn = response.get('cloud_audience_urn', "urn:federation:MicrosoftOnline")
 
         self._log_parsed_response()
 

adal/wstrust_request.py

@@ -41,10 +41,10 @@
 
 class WSTrustRequest(object):
 
-    def __init__(self, call_context, watrust_endpoint_url, applies_to, wstrust_endpoint_version):
+    def __init__(self, call_context, wstrust_endpoint_url, applies_to, wstrust_endpoint_version):
         self._log = log.Logger('WSTrustRequest', call_context['log_context'])
         self._call_context = call_context
-        self._wstrust_endpoint_url = watrust_endpoint_url
+        self._wstrust_endpoint_url = wstrust_endpoint_url
         self._applies_to = applies_to
         self._wstrust_endpoint_version = wstrust_endpoint_version