expephalon-kumisms/totp/management/commands/disabletotp.py

23 lines
788 B
Python
Raw Normal View History

from django.core.management.base import BaseCommand, CommandError
from django.contrib.auth import get_user_model
from totp.models import TOTPUser
from totp.otp import TOTP
class Command(BaseCommand):
help = 'Disables TOTP for the specified user (identified by username)'
def add_arguments(self, parser):
parser.add_argument('user', type=str)
def handle(self, *args, **options):
try:
user = get_user_model().objects.get(username=options["user"])
except get_user_model().DoesNotExist:
raise ValueError(f"User {options['user']} does not exist")
try:
TOTPUser.objects.get(user=user).delete()
except TOTPUser.DoesNotExist:
raise ValueError(f"TOTP not enabled for user {options['user']}")