diff --git a/src/wagtail_2fa/middleware.py b/src/wagtail_2fa/middleware.py index 613c0de..22d9f7e 100644 --- a/src/wagtail_2fa/middleware.py +++ b/src/wagtail_2fa/middleware.py @@ -96,7 +96,7 @@ def process_request(self, request): # Add an attribute to the user so we can easily determine if 2FA should # be enabled for them. - request.user.enable_2fa = request.user.has_perms(["wagtailadmin.enable_2fa"]) + request.user.enable_2fa = request.user.has_perms(["wagtail_2fa.enable_2fa"]) return result @@ -107,7 +107,7 @@ def _require_verified_user(self, request): # 2FA disabled. user_has_device = django_otp.user_has_device(request.user, confirmed=True) if not user_has_device and not request.user.has_perms( - ["wagtailadmin.enable_2fa"] + ["wagtail_2fa.enable_2fa"] ): return False diff --git a/src/wagtail_2fa/migrations/0002_custom_permission.py b/src/wagtail_2fa/migrations/0002_custom_permission.py new file mode 100644 index 0000000..42126b1 --- /dev/null +++ b/src/wagtail_2fa/migrations/0002_custom_permission.py @@ -0,0 +1,43 @@ +from django.db import migrations + + +def create_2fa_permissions(apps, schema_editor): + ContentType = apps.get_model('contenttypes.ContentType') + Permission = apps.get_model('auth.Permission') + + wagtail_2fa_content_type, created = ContentType.objects.get_or_create( + app_label='wagtail_2fa', + model='admin' + ) + + # Create 2FA permission + enable_2fa_permission, created = Permission.objects.get_or_create( + content_type=wagtail_2fa_content_type, + codename='enable_2fa', + name='Enable 2FA' + ) + + +def remove_2fa_permissions(apps, schema_editor): + """Reverse the above additions of permissions.""" + ContentType = apps.get_model('contenttypes.ContentType') + Permission = apps.get_model('auth.Permission') + wagtail_2fa_content_type = ContentType.objects.get( + app_label='wagtail_2fa', + model='admin', + ) + + # This also removes the permission from all groups + Permission.objects.filter( + content_type=wagtail_2fa_content_type, + codename='enable_2fa', + ).delete() + + +class Migration(migrations.Migration): + + dependencies = [] + + operations = [ + migrations.RunPython(create_2fa_permissions, remove_2fa_permissions), + ] diff --git a/src/wagtail_2fa/wagtail_hooks.py b/src/wagtail_2fa/wagtail_hooks.py index c7ec351..6cdc970 100644 --- a/src/wagtail_2fa/wagtail_hooks.py +++ b/src/wagtail_2fa/wagtail_hooks.py @@ -81,7 +81,7 @@ def register_user_listing_buttons(context, user): def register_2fa_permission(): if "wagtail_2fa.middleware.VerifyUserPermissionsMiddleware" in settings.MIDDLEWARE: return Permission.objects.filter( - content_type__app_label="wagtailadmin", codename="enable_2fa" + content_type__app_label="wagtail_2fa", codename="enable_2fa" ) return Permission.objects.none()