mirror of
				https://github.com/inventree/InvenTree.git
				synced 2025-10-30 04:35:42 +00:00 
			
		
		
		
	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
This commit is contained in:
		| @@ -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 | ||||
|   | ||||
| @@ -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])) | ||||
|   | ||||
		Reference in New Issue
	
	Block a user