From 76dfd62f8493366d92a9be75fdde14be23f18cdb Mon Sep 17 00:00:00 2001
From: Matthias Mair <code@mjmair.com>
Date: Wed, 20 Aug 2025 13:21:03 +0200
Subject: [PATCH] fix: remove mfa task (#10200)

* fix remove mfa task

* add test for command

* clean up after schema test

* add assert to esure authenticators are really present/removed

* simplify handler
---
 .../management/commands/remove_mfa.py         | 29 +++++++++------
 .../InvenTree/InvenTree/test_commands.py      | 37 +++++++++++++++++++
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/src/backend/InvenTree/InvenTree/management/commands/remove_mfa.py b/src/backend/InvenTree/InvenTree/management/commands/remove_mfa.py
index 60a7cac51f..1acceeb1ed 100644
--- a/src/backend/InvenTree/InvenTree/management/commands/remove_mfa.py
+++ b/src/backend/InvenTree/InvenTree/management/commands/remove_mfa.py
@@ -3,6 +3,10 @@
 from django.contrib.auth import get_user_model
 from django.core.management.base import BaseCommand
 
+import structlog
+
+logger = structlog.get_logger('inventree')
+
 
 class Command(BaseCommand):
     """Remove MFA for a user."""
@@ -11,12 +15,8 @@ class Command(BaseCommand):
         """Add the arguments."""
         parser.add_argument('mail', type=str)
 
-    def handle(self, *args, **kwargs):
+    def handle(self, *args, mail, **kwargs):
         """Remove MFA for the supplied user (by mail)."""
-        # general settings
-        mail = kwargs.get('mail')
-        if not mail:
-            raise KeyError('A mail is required')
         user = get_user_model()
         mfa_user = [
             *set(
@@ -26,13 +26,18 @@ class Command(BaseCommand):
         ]
 
         if len(mfa_user) == 0:
-            print('No user with this mail associated')
+            logger.warning('No user with this mail associated')
         elif len(mfa_user) > 1:
-            print('More than one user found with this mail')
+            logger.error('More than one user found with this mail')
         else:
             # and clean out all MFA methods
-            # backup codes
-            mfa_user[0].staticdevice_set.all().delete()
-            # TOTP tokens
-            mfa_user[0].totpdevice_set.all().delete()
-            print(f'Removed all MFA methods for user {mfa_user[0]!s}')
+            auths = mfa_user[0].authenticator_set.all()
+            length = len(auths)
+            auths.delete()
+
+            # log the result
+            msg = f'Removed all ({length}) MFA methods for user {mfa_user[0]!s}'
+            logger.info(msg)
+            print(msg)
+            return 'done'
+        return False
diff --git a/src/backend/InvenTree/InvenTree/test_commands.py b/src/backend/InvenTree/InvenTree/test_commands.py
index c1d5e56761..43e143e91b 100644
--- a/src/backend/InvenTree/InvenTree/test_commands.py
+++ b/src/backend/InvenTree/InvenTree/test_commands.py
@@ -1,5 +1,8 @@
 """Tests for custom InvenTree management commands."""
 
+from pathlib import Path
+
+from django.contrib.auth.models import User
 from django.core.management import call_command
 from django.test import TestCase
 
@@ -11,3 +14,37 @@ class CommandTestCase(TestCase):
         """Test the schema generation command."""
         output = call_command('schema', file='schema.yml', verbosity=0)
         self.assertEqual(output, 'done')
+
+        Path('schema.yml').unlink()  # cleanup
+
+    def test_remove_mfa(self):
+        """Test the remove_mfa command."""
+        # missing arg
+        with self.assertRaises(Exception) as cm:
+            call_command('remove_mfa', verbosity=0)
+        self.assertEqual(
+            'Error: the following arguments are required: mail', str(cm.exception)
+        )
+
+        # no user
+        with self.assertLogs('inventree') as cm:
+            self.assertFalse(
+                call_command('remove_mfa', 'admin@example.org', verbosity=0)
+            )
+        self.assertIn('No user with this mail associated', str(cm[1]))
+
+        # correct removal
+        my_admin = User.objects.create_user(username='admin', email='admin@example.org')
+        my_admin.authenticator_set.create(type='TOTP', data={})
+        self.assertEqual(my_admin.authenticator_set.all().count(), 1)
+        output = call_command('remove_mfa', 'admin@example.org', verbosity=0)
+        self.assertEqual(output, 'done')
+        self.assertEqual(my_admin.authenticator_set.all().count(), 0)
+
+        # two users with same email
+        User.objects.create_user(username='admin2', email='admin@example.org')
+        with self.assertLogs('inventree') as cm:
+            self.assertFalse(
+                call_command('remove_mfa', 'admin@example.org', verbosity=0)
+            )
+        self.assertIn('More than one user found with this mail', str(cm[1]))