Update token endpoint

It needs handled differently depending on the device grant type or not
it also needs to be rate limited to adhrere to the polling section in the spec
so a device can't spam the token endpoint

Author: duzumaki

oauth2_provider/views/base.py

@@ -3,6 +3,7 @@
 import logging
 from urllib.parse import parse_qsl, urlencode, urlparse
 
+from django import http
 from django.contrib.auth.mixins import LoginRequiredMixin
 from django.contrib.auth.views import redirect_to_login
 from django.http import HttpResponse
@@ -12,6 +13,13 @@
 from django.views.decorators.csrf import csrf_exempt
 from django.views.decorators.debug import sensitive_post_parameters
 from django.views.generic import FormView, View
+from django_ratelimit.decorators import ratelimit
+from oauthlib.oauth2.rfc8628.errors import (
+    AccessDenied,
+    AuthorizationPendingError,
+)
+
+from oauth2_provider.models import Device
 
 from ..compat import login_not_required
 from ..exceptions import OAuthToolkitError
@@ -290,10 +298,13 @@ class TokenView(OAuthLibMixin, View):
     * Authorization code
     * Password
     * Client credentials
+    * Device code flow (specifically for the device polling stage)
     """
 
     @method_decorator(sensitive_post_parameters("password", "client_secret"))
-    def post(self, request, *args, **kwargs):
+    def authorization_flow_token_response(
+        self, request: http.HttpRequest, *args, **kwargs
+    ) -> http.HttpResponse:
         url, headers, body, status = self.create_token_response(request)
         if status == 200:
             access_token = json.loads(body).get("access_token")
@@ -307,6 +318,38 @@ def post(self, request, *args, **kwargs):
             response[k] = v
         return response
 
+    @method_decorator(ratelimit(key="ip", rate=f"1/{oauth2_settings.DEVICE_FLOW_INTERVAL}"))
+    def device_flow_token_response(
+        self, request: http.HttpRequest, device_code: str, *args, **kwargs
+    ) -> http.HttpResponse:
+        device = Device.objects.get(device_code=device_code)
+
+        if device.status == device.AUTHORIZATION_PENDING:
+            raise AuthorizationPendingError
+
+        if device.status == device.DENIED:
+            raise AccessDenied
+
+        url, headers, body, status = self.create_token_response(request)
+
+        if status != 200:
+            return http.JsonResponse(data=json.loads(body), status=status)
+
+        response = http.JsonResponse(data=json.loads(body), status=status)
+
+        for k, v in headers.items():
+            response[k] = v
+
+        device.status = device.EXPIRED
+        device.save(update_fields=["status"])
+        return response
+
+    def post(self, request: http.HttpRequest, *args, **kwargs) -> http.HttpResponse:
+        params = request.POST
+        if params.get("grant_type") == "urn:ietf:params:oauth:grant-type:device_code":
+            return self.device_flow_token_response(request, params["device_code"])
+        return self.authorization_flow_token_response(request)
+
 
 @method_decorator(csrf_exempt, name="dispatch")
 @method_decorator(login_not_required, name="dispatch")