Add management command from #189 (comment)

allauth_2fa/management/commands/allauth_2fa_migrate.py

@@ -0,0 +1,37 @@
+import base64
+
+from allauth.mfa.adapter import get_adapter
+from allauth.mfa.models import Authenticator
+from django.core.management.base import BaseCommand
+from django_otp.plugins.otp_static.models import StaticDevice
+from django_otp.plugins.otp_totp.models import TOTPDevice
+
+
+class Command(BaseCommand):
+    def handle(self, **options):
+        adapter = get_adapter()
+        authenticators = []
+        for totp in TOTPDevice.objects.filter(confirmed=True).iterator():
+            recovery_codes = set()
+            for sdevice in StaticDevice.objects.filter(
+                confirmed=True, user_id=totp.user_id
+            ).iterator():
+                recovery_codes.update(sdevice.token_set.values_list("token", flat=True))
+            secret = base64.b32encode(bytes.fromhex(totp.key)).decode("ascii")
+            totp_authenticator = Authenticator(
+                user_id=totp.user_id,
+                type=Authenticator.Type.TOTP,
+                data={"secret": adapter.encrypt(secret)},
+            )
+            authenticators.append(totp_authenticator)
+            authenticators.append(
+                Authenticator(
+                    user_id=totp.user_id,
+                    type=Authenticator.Type.RECOVERY_CODES,
+                    data={
+                        "migrated_codes": [adapter.encrypt(c) for c in recovery_codes],
+                    },
+                )
+            )
+        Authenticator.objects.bulk_create(authenticators)
+        self.stdout.write(f"Created {len(authenticators)} Authenticators")

tests/settings.py

@@ -46,6 +46,7 @@
     # Enable allauth.
     "allauth",
     "allauth.account",
+    "allauth.mfa",  # For testing the migration.
     # Required to render the default template for 'account_login'.
     "allauth.socialaccount",
     # Configure the django-otp package.

tests/test_allauth_2fa.py

@@ -9,6 +9,7 @@
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import AbstractUser
+from django.core.management import call_command
 from django.forms import BaseForm
 from django.test import override_settings
 from django.urls import reverse
@@ -48,6 +49,13 @@ def pytest_generate_tests(metafunc):
         metafunc.parametrize("adapter", ADAPTER_CLASSES, indirect=True)
 
 
+def create_totp_and_static(user: AbstractUser) -> tuple[TOTPDevice, StaticDevice]:
+    totp_model = user.totpdevice_set.create()
+    static_model = user.staticdevice_set.create()
+    static_model.token_set.create(token=StaticToken.random_token())
+    return totp_model, static_model
+
+
 @pytest.fixture(autouse=True)
 def adapter(request, settings):
     settings.ACCOUNT_ADAPTER = request.param
@@ -65,12 +73,11 @@ def john() -> AbstractUser:
 
 @pytest.fixture()
 def john_with_totp(john: AbstractUser) -> tuple[AbstractUser, TOTPDevice, StaticDevice]:
-    totp_model = john.totpdevice_set.create()
-    static_model = john.staticdevice_set.create()
-    static_model.token_set.create(token=StaticToken.random_token())
+    totp_model, static_model = create_totp_and_static(john)
     return john, totp_model, static_model
 
 
+
 @pytest.fixture()
 def user_logged_in_count(request) -> Callable[[], int]:
     login_callback = Mock()
@@ -404,3 +411,13 @@ def test_view_missing_attribute(request, view_cls) -> None:
 
     # Ensure the function doesn't fail when the attribute is missing.
     assert OTPAdapter().get_2fa_authenticate_url(request) is not None
+
+
+def test_migration_management_command():
+    from allauth.mfa.models import Authenticator
+    for x in range(10):
+        user = get_user_model().objects.create(username=f"user{x}")
+        create_totp_and_static(user)
+    call_command("allauth_2fa_migrate")
+    assert Authenticator.objects.filter(type=Authenticator.Type.RECOVERY_CODES).count() == 10
+    assert Authenticator.objects.filter(type=Authenticator.Type.TOTP).count() == 10