From bd1e019a4201a7c33277b2b2e5b4a395de3d5e77 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@ns1.com>
Date: Mon, 28 Jun 2021 16:19:02 -0400
Subject: [PATCH] Clean up token-based authentication for GraphQL

---
 netbox/netbox/graphql/views.py      | 40 ++++++++++++++++++++---------
 netbox/netbox/tests/test_graphql.py |  2 +-
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/netbox/netbox/graphql/views.py b/netbox/netbox/graphql/views.py
index 1cde56cd6..047f7bab3 100644
--- a/netbox/netbox/graphql/views.py
+++ b/netbox/netbox/graphql/views.py
@@ -1,20 +1,36 @@
+from django.conf import settings
+from django.contrib.auth.views import redirect_to_login
+from django.http import HttpResponseForbidden
+from django.urls import reverse
 from graphene_django.views import GraphQLView as GraphQLView_
-from rest_framework.decorators import authentication_classes, permission_classes, api_view
-from rest_framework.permissions import IsAuthenticated
-from rest_framework.settings import api_settings
+from rest_framework.exceptions import AuthenticationFailed
+
+from netbox.api.authentication import TokenAuthentication
 
 
 class GraphQLView(GraphQLView_):
     """
-    Extends grpahene_django's GraphQLView to support DRF's token-based authentication.
+    Extends graphene_django's GraphQLView to support DRF's token-based authentication.
     """
-    @classmethod
-    def as_view(cls, *args, **kwargs):
-        view = super(GraphQLView, cls).as_view(*args, **kwargs)
+    def dispatch(self, request, *args, **kwargs):
 
-        # Apply DRF permission and authentication classes
-        view = permission_classes((IsAuthenticated,))(view)
-        view = authentication_classes(api_settings.DEFAULT_AUTHENTICATION_CLASSES)(view)
-        view = api_view(['GET', 'POST'])(view)
+        # Attempt to authenticate the user using a DRF token, if provided
+        if not request.user.is_authenticated:
+            authenticator = TokenAuthentication()
+            try:
+                auth_info = authenticator.authenticate(request)
+                if auth_info is not None:
+                    request.user = auth_info[0]  # User object
+            except AuthenticationFailed as exc:
+                return HttpResponseForbidden(exc.detail)
 
-        return view
+        # Enforce LOGIN_REQUIRED
+        if settings.LOGIN_REQUIRED and not request.user.is_authenticated:
+
+            # If this is a human user, send a redirect to the login page
+            if self.request_wants_html(request):
+                return redirect_to_login(reverse('graphql'))
+
+            return HttpResponseForbidden("No credentials provided.")
+
+        return super().dispatch(request, *args, **kwargs)
diff --git a/netbox/netbox/tests/test_graphql.py b/netbox/netbox/tests/test_graphql.py
index dd43bbbdd..483c125a2 100644
--- a/netbox/netbox/tests/test_graphql.py
+++ b/netbox/netbox/tests/test_graphql.py
@@ -24,4 +24,4 @@ class GraphQLTestCase(TestCase):
         self.client.logout()
         response = self.client.get(url, **header)
         with disable_warnings('django.request'):
-            self.assertHttpStatus(response, 302)
+            self.assertHttpStatus(response, 302)  # Redirect to login page