From f446d53a4eab227c7db2a997805e6830c6383bf6 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Sat, 19 Oct 2024 22:55:34 -0700
Subject: [PATCH] Backend tests

---
 src/documents/tests/test_api_permissions.py | 120 ++++++++++++++++++++
 src/paperless/views.py                      |   3 +-
 2 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/src/documents/tests/test_api_permissions.py b/src/documents/tests/test_api_permissions.py
index 7708b8541..2f14ab951 100644
--- a/src/documents/tests/test_api_permissions.py
+++ b/src/documents/tests/test_api_permissions.py
@@ -1,5 +1,7 @@
 import json
+from unittest import mock
 
+from allauth.mfa.models import Authenticator
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
@@ -601,6 +603,51 @@ class TestApiUser(DirectoriesMixin, APITestCase):
         self.assertEqual(returned_user2.first_name, "Updated Name 2")
         self.assertNotEqual(returned_user2.password, initial_password)
 
+    def test_deactivate_totp(self):
+        """
+        GIVEN:
+            - Existing user account with TOTP enabled
+        WHEN:
+            - API request by a superuser is made to deactivate TOTP
+        THEN:
+            - TOTP is deactivated
+        """
+
+        user1 = User.objects.create(
+            username="testuser",
+            password="test",
+            first_name="Test",
+            last_name="User",
+        )
+        Authenticator.objects.create(
+            user=user1,
+            type=Authenticator.Type.TOTP,
+            data={},
+        )
+
+        response = self.client.post(
+            f"{self.ENDPOINT}{user1.pk}/deactivate_totp/",
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(Authenticator.objects.filter(user=user1).count(), 0)
+
+        regular_user = User.objects.create_user(username="regular_user")
+        regular_user.user_permissions.add(
+            *Permission.objects.all(),
+        )
+        self.client.force_authenticate(regular_user)
+        Authenticator.objects.create(
+            user=user1,
+            type=Authenticator.Type.TOTP,
+            data={},
+        )
+
+        response = self.client.post(
+            f"{self.ENDPOINT}{user1.pk}/deactivate_totp/",
+        )
+        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
 
 class TestApiGroup(DirectoriesMixin, APITestCase):
     ENDPOINT = "/api/groups/"
@@ -1012,3 +1059,76 @@ class TestBulkEditObjectPermissions(APITestCase):
         )
 
         self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+
+class TestApiTOTPViews(APITestCase):
+    ENDPOINT = "/api/profile/totp/"
+
+    def setUp(self):
+        super().setUp()
+
+        self.user = User.objects.create_superuser(username="temp_admin")
+        self.client.force_authenticate(user=self.user)
+
+    def test_get_totp(self):
+        """
+        GIVEN:
+            - Existing user account
+        WHEN:
+            - API request is made to TOTP endpoint
+        THEN:
+            - TOTP is generated
+        """
+        response = self.client.get(
+            self.ENDPOINT,
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertIn("qr_svg", response.data)
+        self.assertIn("secret", response.data)
+
+    @mock.patch("allauth.mfa.totp.internal.auth.validate_totp_code")
+    def test_activate_totp(self, mock_validate_totp_code):
+        """
+        GIVEN:
+            - Existing user account
+        WHEN:
+            - API request is made to activate TOTP
+        THEN:
+            - TOTP is activated, recovery codes are returned
+        """
+        mock_validate_totp_code.return_value = True
+
+        response = self.client.post(
+            self.ENDPOINT,
+            data={
+                "secret": "123",
+                "code": "456",
+            },
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertTrue(Authenticator.objects.filter(user=self.user).exists())
+        self.assertIn("recovery_codes", response.data)
+
+    def test_deactivate_totp(self):
+        """
+        GIVEN:
+            - Existing user account with TOTP enabled
+        WHEN:
+            - API request is made to deactivate TOTP
+        THEN:
+            - TOTP is deactivated
+        """
+        Authenticator.objects.create(
+            user=self.user,
+            type=Authenticator.Type.TOTP,
+            data={},
+        )
+
+        response = self.client.delete(
+            self.ENDPOINT,
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(Authenticator.objects.filter(user=self.user).count(), 0)
diff --git a/src/paperless/views.py b/src/paperless/views.py
index ea9d0b89e..a7dff8282 100644
--- a/src/paperless/views.py
+++ b/src/paperless/views.py
@@ -14,6 +14,7 @@ from django.contrib.auth.models import User
 from django.db.models.functions import Lower
 from django.http import HttpResponse
 from django.http import HttpResponseBadRequest
+from django.http import HttpResponseForbidden
 from django.views.generic import View
 from django_filters.rest_framework import DjangoFilterBackend
 from rest_framework.authtoken.models import Token
@@ -112,7 +113,7 @@ class UserViewSet(ModelViewSet):
         request_user = request.user
         user = User.objects.get(pk=pk)
         if not request_user.is_superuser and request_user != user:
-            return HttpResponseBadRequest(
+            return HttpResponseForbidden(
                 "You do not have permission to deactivate TOTP for this user",
             )
         try: