Add is_staff to role mappings, and the logic to go with it

2 files changed, 63 insertions, 11 deletions

diff --git a/pydis_site/apps/home/signals.py b/pydis_site/apps/home/signals.py
index 4bf50b78..76d29631 100644
--- a/pydis_site/apps/home/signals.py
+++ b/pydis_site/apps/home/signals.py
@@ -1,3 +1,4 @@
+from contextlib import suppress
 from typing import List, Optional, Type
 
 from allauth.account.signals import user_logged_in
@@ -8,7 +9,7 @@ from allauth.socialaccount.signals import (
     pre_social_login, social_account_added, social_account_removed,
     social_account_updated)
 from django.contrib.auth.models import Group, User as DjangoUser
-from django.db.models.signals import post_save, pre_delete, pre_save
+from django.db.models.signals import post_delete, post_save, pre_save
 
 from pydis_site.apps.api.models import User as DiscordUser
 from pydis_site.apps.staff.models import RoleMapping
@@ -37,7 +38,7 @@ class AllauthSignalListener:
     def __init__(self):
         post_save.connect(self.user_model_updated, sender=DiscordUser)
 
-        pre_delete.connect(self.mapping_model_deleted, sender=RoleMapping)
+        post_delete.connect(self.mapping_model_deleted, sender=RoleMapping)
         pre_save.connect(self.mapping_model_updated, sender=RoleMapping)
 
         pre_social_login.connect(self.social_account_updated)
@@ -133,13 +134,29 @@ class AllauthSignalListener:
         Processes deletion signals from the RoleMapping model, removing perms from users.
 
         We need to do this to ensure that users aren't left with permissions groups that
-        they shouldn't have assigned to them when a RoleMapping is deleted from the database.
+        they shouldn't have assigned to them when a RoleMapping is deleted from the database,
+        and to remove their staff status if they should no longer have it.
         """
         instance: RoleMapping = kwargs["instance"]
 
         for user in instance.group.user_set.all():
+            # Firstly, remove their related user group
             user.groups.remove(instance.group)
 
+            with suppress(SocialAccount.DoesNotExist, DiscordUser.DoesNotExist):
+                # If we get either exception, then the user could not have been assigned staff
+                # with our system in the first place.
+
+                social_account = SocialAccount.objects.get(user=user, provider=DiscordProvider.id)
+                discord_user = DiscordUser.objects.get(id=int(social_account.uid))
+
+                mappings = RoleMapping.objects.filter(role__in=discord_user.roles.all()).all()
+                is_staff = any(m.is_staff for m in mappings)
+
+                if user.is_staff != is_staff:
+                    user.is_staff = is_staff
+                    user.save(update_fields=("is_staff", ))
+
     def mapping_model_updated(self, sender: Type[RoleMapping], **kwargs) -> None:
         """
         Processes update signals from the RoleMapping model.
@@ -174,6 +191,21 @@ class AllauthSignalListener:
         for account in accounts:
             account.user.groups.add(instance.group)
 
+            if instance.is_staff and not account.user.is_staff:
+                account.user.is_staff = True
+                account.user.save(update_fields=("is_staff", ))
+            else:
+                discord_user = DiscordUser.objects.get(id=int(account.uid))
+
+                mappings = RoleMapping.objects.filter(
+                    role__in=discord_user.roles.all()
+                ).exclude(id=instance.id).all()
+                is_staff = any(m.is_staff for m in mappings)
+
+                if account.user.is_staff != is_staff:
+                    account.user.is_staff = is_staff
+                    account.user.save(update_fields=("is_staff",))
+
     def user_model_updated(self, sender: Type[DiscordUser], **kwargs) -> None:
         """
         Processes update signals from the Discord User model, assigning perms as required.
@@ -235,22 +267,33 @@ class AllauthSignalListener:
 
         if deletion:
             # They've unlinked Discord or left the server, so we have to remove their groups
-
-            if not current_groups:
-                return  # They have no groups anyway, no point in processing
-
-            account.user.groups.remove(
-                *(mapping.group for mapping in mappings)
-            )
+            # and their staff status
+
+            if current_groups:
+                # They do have groups, so let's remove them
+                account.user.groups.remove(
+                    *(mapping.group for mapping in mappings)
+                )
+
+            if account.user.is_staff:
+                # They're marked as a staff user and they shouldn't be, so let's fix that
+                account.user.is_staff = False
+                account.user.save(update_fields=("is_staff", ))
         else:
             new_groups = []
+            is_staff = False
 
             for role in user.roles.all():
                 try:
-                    new_groups.append(mappings.get(role=role).group)
+                    mapping = mappings.get(role=role)
                 except RoleMapping.DoesNotExist:
                     continue  # No mapping exists
 
+                new_groups.append(mapping.group)
+
+                if mapping.is_staff:
+                    is_staff = True
+
             account.user.groups.add(
                 *[group for group in new_groups if group not in current_groups]
             )
@@ -258,3 +301,7 @@ class AllauthSignalListener:
             account.user.groups.remove(
                 *[mapping.group for mapping in mappings if mapping.group not in new_groups]
             )
+
+            if account.user.is_staff != is_staff:
+                account.user.is_staff = is_staff
+                account.user.save(update_fields=("is_staff", ))
diff --git a/pydis_site/apps/staff/models/role_mapping.py b/pydis_site/apps/staff/models/role_mapping.py
index 10c09cf1..dff8081a 100644
--- a/pydis_site/apps/staff/models/role_mapping.py
+++ b/pydis_site/apps/staff/models/role_mapping.py
@@ -21,6 +21,11 @@ class RoleMapping(models.Model):
         unique=True,  # Unique in order to simplify group assignment logic
     )
 
+    is_staff = models.BooleanField(
+        help_text="Whether this role mapping related to a Django staff group",
+        default=False
+    )
+
     def __str__(self):
         """Returns the mapping, for display purposes."""
         return f"@{self.role.name} -> {self.group.name}"