Enhancement: require totp code for obtain auth token (#8936)

src/paperless/serialisers.py

@@ -1,11 +1,14 @@
 import logging
 
 from allauth.mfa.adapter import get_adapter as get_mfa_adapter
+from allauth.mfa.models import Authenticator
+from allauth.mfa.totp.internal.auth import TOTP
 from allauth.socialaccount.models import SocialAccount
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
 from rest_framework import serializers
+from rest_framework.authtoken.serializers import AuthTokenSerializer
 
 from paperless.models import ApplicationConfiguration
 
@@ -24,6 +27,36 @@ class ObfuscatedUserPasswordField(serializers.Field):
         return data
 
 
+class PaperlessAuthTokenSerializer(AuthTokenSerializer):
+    code = serializers.CharField(
+        label="MFA Code",
+        write_only=True,
+        required=False,
+    )
+
+    def validate(self, attrs):
+        attrs = super().validate(attrs)
+        user = attrs.get("user")
+        code = attrs.get("code")
+        mfa_adapter = get_mfa_adapter()
+        if mfa_adapter.is_mfa_enabled(user):
+            if not code:
+                raise serializers.ValidationError(
+                    "MFA code is required",
+                )
+            authenticator = Authenticator.objects.get(
+                user=user,
+                type=Authenticator.Type.TOTP,
+            )
+            if not TOTP(instance=authenticator).validate_code(
+                code,
+            ):
+                raise serializers.ValidationError(
+                    "Invalid MFA code",
+                )
+        return attrs
+
+
 class UserSerializer(serializers.ModelSerializer):
     password = ObfuscatedUserPasswordField(required=False)
     user_permissions = serializers.SlugRelatedField(

src/paperless/urls.py

@@ -14,7 +14,6 @@ from django.utils.translation import gettext_lazy as _
 from django.views.decorators.csrf import ensure_csrf_cookie
 from django.views.generic import RedirectView
 from django.views.static import serve
-from rest_framework.authtoken import views
 from rest_framework.routers import DefaultRouter
 
 from documents.views import BulkDownloadView
@@ -50,6 +49,7 @@ from paperless.views import DisconnectSocialAccountView
 from paperless.views import FaviconView
 from paperless.views import GenerateAuthTokenView
 from paperless.views import GroupViewSet
+from paperless.views import PaperlessObtainAuthTokenView
 from paperless.views import ProfileView
 from paperless.views import SocialAccountProvidersView
 from paperless.views import TOTPView
@@ -157,7 +157,7 @@ urlpatterns = [
                 ),
                 path(
                     "token/",
-                    views.obtain_auth_token,
+                    PaperlessObtainAuthTokenView.as_view(),
                 ),
                 re_path(
                     "^profile/",

src/paperless/views.py

@@ -19,6 +19,7 @@ from django.http import HttpResponseNotFound
 from django.views.generic import View
 from django_filters.rest_framework import DjangoFilterBackend
 from rest_framework.authtoken.models import Token
+from rest_framework.authtoken.views import ObtainAuthToken
 from rest_framework.decorators import action
 from rest_framework.filters import OrderingFilter
 from rest_framework.generics import GenericAPIView
@@ -35,10 +36,15 @@ from paperless.filters import UserFilterSet
 from paperless.models import ApplicationConfiguration
 from paperless.serialisers import ApplicationConfigurationSerializer
 from paperless.serialisers import GroupSerializer
+from paperless.serialisers import PaperlessAuthTokenSerializer
 from paperless.serialisers import ProfileSerializer
 from paperless.serialisers import UserSerializer
 
 
+class PaperlessObtainAuthTokenView(ObtainAuthToken):
+    serializer_class = PaperlessAuthTokenSerializer
+
+
 class StandardPagination(PageNumberPagination):
     page_size = 25
     page_size_query_param = "page_size"