mirror of
https://github.com/inventree/InvenTree.git
synced 2025-09-13 14:11:37 +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.contrib.auth import get_user_model
|
||||||
from django.core.management.base import BaseCommand
|
from django.core.management.base import BaseCommand
|
||||||
|
|
||||||
|
import structlog
|
||||||
|
|
||||||
|
logger = structlog.get_logger('inventree')
|
||||||
|
|
||||||
|
|
||||||
class Command(BaseCommand):
|
class Command(BaseCommand):
|
||||||
"""Remove MFA for a user."""
|
"""Remove MFA for a user."""
|
||||||
@@ -11,12 +15,8 @@ class Command(BaseCommand):
|
|||||||
"""Add the arguments."""
|
"""Add the arguments."""
|
||||||
parser.add_argument('mail', type=str)
|
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)."""
|
"""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()
|
user = get_user_model()
|
||||||
mfa_user = [
|
mfa_user = [
|
||||||
*set(
|
*set(
|
||||||
@@ -26,13 +26,18 @@ class Command(BaseCommand):
|
|||||||
]
|
]
|
||||||
|
|
||||||
if len(mfa_user) == 0:
|
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:
|
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:
|
else:
|
||||||
# and clean out all MFA methods
|
# and clean out all MFA methods
|
||||||
# backup codes
|
auths = mfa_user[0].authenticator_set.all()
|
||||||
mfa_user[0].staticdevice_set.all().delete()
|
length = len(auths)
|
||||||
# TOTP tokens
|
auths.delete()
|
||||||
mfa_user[0].totpdevice_set.all().delete()
|
|
||||||
print(f'Removed all MFA methods for user {mfa_user[0]!s}')
|
# 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."""
|
"""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.core.management import call_command
|
||||||
from django.test import TestCase
|
from django.test import TestCase
|
||||||
|
|
||||||
@@ -11,3 +14,37 @@ class CommandTestCase(TestCase):
|
|||||||
"""Test the schema generation command."""
|
"""Test the schema generation command."""
|
||||||
output = call_command('schema', file='schema.yml', verbosity=0)
|
output = call_command('schema', file='schema.yml', verbosity=0)
|
||||||
self.assertEqual(output, 'done')
|
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