Consolidate request handling logic into helper method.

The _ms_request conslidates the error handling logic into
a single method. This should make overriding methods easier
for end users.

Author: tim-schilling

django_auth_adfs/backend.py

@@ -15,19 +15,22 @@
 
 
 class AdfsBaseBackend(ModelBackend):
-    def exchange_auth_code(self, authorization_code, request):
-        logger.debug("Received authorization code: %s", authorization_code)
-        data = {
-            'grant_type': 'authorization_code',
-            'client_id': settings.CLIENT_ID,
-            'redirect_uri': provider_config.redirect_uri(request),
-            'code': authorization_code,
-        }
-        if settings.CLIENT_SECRET:
-            data['client_secret'] = settings.CLIENT_SECRET
 
-        logger.debug("Getting access token at: %s", provider_config.token_endpoint)
-        response = provider_config.session.post(provider_config.token_endpoint, data, timeout=settings.TIMEOUT)
+    def _ms_request(self, action, url, data=None, **kwargs):
+        """
+        Make a Microsoft Entra/GraphQL request
+
+
+        Args:
+            action (callable): The callable for making a request.
+            url (str): The URL the request should be sent to.
+            data (dict): Optional dictionary of data to be sent in the request.
+
+        Returns:
+            response: The response from the server. If it's not a 200, a
+                      PermissionDenied is raised.
+        """
+        response = action(url, data=data, timeout=settings.TIMEOUT, **kwargs)
         # 200 = valid token received
         # 400 = 'something' is wrong in our request
         if response.status_code == 400:
@@ -39,7 +42,21 @@ def exchange_auth_code(self, authorization_code, request):
         if response.status_code != 200:
             logger.error("Unexpected ADFS response: %s", response.content.decode())
             raise PermissionDenied
+        return response
+
+    def exchange_auth_code(self, authorization_code, request):
+        logger.debug("Received authorization code: %s", authorization_code)
+        data = {
+            'grant_type': 'authorization_code',
+            'client_id': settings.CLIENT_ID,
+            'redirect_uri': provider_config.redirect_uri(request),
+            'code': authorization_code,
+        }
+        if settings.CLIENT_SECRET:
+            data['client_secret'] = settings.CLIENT_SECRET
 
+        logger.debug("Getting access token at: %s", provider_config.token_endpoint)
+        response = self._ms_request(provider_config.session.post, provider_config.token_endpoint, data)
         adfs_response = response.json()
         return adfs_response
 
@@ -66,17 +83,7 @@ def get_obo_access_token(self, access_token):
         else:
             data["resource"] = 'https://graph.microsoft.com'
 
-        response = provider_config.session.get(provider_config.token_endpoint, data=data, timeout=settings.TIMEOUT)
-        # 200 = valid token received
-        # 400 = 'something' is wrong in our request
-        if response.status_code == 400:
-            logger.error("ADFS server returned an error: %s", response.json()["error_description"])
-            raise PermissionDenied
-
-        if response.status_code != 200:
-            logger.error("Unexpected ADFS response: %s", response.content.decode())
-            raise PermissionDenied
-
+        response = self._ms_request(provider_config.session.get, provider_config.token_endpoint, data)
         obo_access_token = response.json()["access_token"]
         logger.debug("Received OBO access token: %s", obo_access_token)
         return obo_access_token
@@ -95,17 +102,11 @@ def get_group_memberships_from_ms_graph(self, obo_access_token):
             provider_config.msgraph_endpoint
         )
         headers = {"Authorization": "Bearer {}".format(obo_access_token)}
-        response = provider_config.session.get(graph_url, headers=headers, timeout=settings.TIMEOUT)
-        # 200 = valid token received
-        # 400 = 'something' is wrong in our request
-        if response.status_code in [400, 401]:
-            logger.error("MS Graph server returned an error: %s", response.json()["message"])
-            raise PermissionDenied
-
-        if response.status_code != 200:
-            logger.error("Unexpected MS Graph response: %s", response.content.decode())
-            raise PermissionDenied
-
+        response = self._ms_request(
+            action=provider_config.session.get,
+            url=graph_url,
+            headers=headers,
+        )
         claim_groups = []
         for group_data in response.json()["value"]:
             if group_data["displayName"] is None:

django_auth_adfs/rest_framework.py

@@ -6,6 +6,8 @@
     BaseAuthentication, get_authorization_header
 )
 
+from django_auth_adfs.exceptions import MFARequired
+
 
 class AdfsAccessTokenAuthentication(BaseAuthentication):
     """
@@ -33,7 +35,10 @@ def authenticate(self, request):
         # Authenticate the user
         # The AdfsAuthCodeBackend authentication backend will notice the "access_token" parameter
         # and skip the request for an access token using the authorization code
-        user = authenticate(access_token=auth[1])
+        try:
+            user = authenticate(access_token=auth[1])
+        except MFARequired as e:
+            raise exceptions.AuthenticationFailed('MFA auth is required.') from e
 
         if user is None:
             raise exceptions.AuthenticationFailed('Invalid access token.')