diff --git a/apps/rbac/models/role.py b/apps/rbac/models/role.py index 9d539ba13..b47bccb5b 100644 --- a/apps/rbac/models/role.py +++ b/apps/rbac/models/role.py @@ -118,6 +118,9 @@ class Role(JMSModel): return self.name return gettext(self.name) + def is_org(self): + return self.scope == const.Scope.org + class SystemRole(Role): objects = SystemRoleManager() diff --git a/apps/rbac/models/rolebinding.py b/apps/rbac/models/rolebinding.py index f7ee3881a..a2ee06022 100644 --- a/apps/rbac/models/rolebinding.py +++ b/apps/rbac/models/rolebinding.py @@ -1,6 +1,7 @@ from django.utils.translation import gettext_lazy as _ from django.db import models from django.db.models import Q +from django.core.exceptions import ValidationError from rest_framework.serializers import ValidationError from common.db.models import JMSModel @@ -67,6 +68,7 @@ class RoleBinding(JMSModel): def save(self, *args, **kwargs): self.scope = self.role.scope + self.clean() return super().save(*args, **kwargs) @classmethod @@ -95,6 +97,9 @@ class RoleBinding(JMSModel): def role_display(self): return self.role.display_name + def is_scope_org(self): + return self.scope == Scope.org + class OrgRoleBindingManager(RoleBindingManager): def get_queryset(self): @@ -147,3 +152,12 @@ class SystemRoleBinding(RoleBinding): def save(self, *args, **kwargs): self.scope = Scope.system return super().save(*args, **kwargs) + + def clean(self): + kwargs = dict(role=self.role, user=self.user, scope=self.scope) + exists = self.__class__.objects.filter(**kwargs).exists() + if exists: + msg = "Duplicate for key 'role_user' of system role binding, {}_{}".format( + self.role.id, self.user.id + ) + raise ValidationError(msg) diff --git a/apps/users/models/user.py b/apps/users/models/user.py index e7834a088..0940d6466 100644 --- a/apps/users/models/user.py +++ b/apps/users/models/user.py @@ -172,6 +172,22 @@ class RoleManager(models.Manager): def __init__(self, user, *args, **kwargs): super().__init__(*args, **kwargs) self.user = user + self.role_binding_cls = self.get_role_binding_cls() + self.role_cls = self.get_role_cls() + + def get_role_binding_cls(self): + from rbac.models import SystemRoleBinding, OrgRoleBinding + if self.scope == 'org': + return OrgRoleBinding + else: + return SystemRoleBinding + + def get_role_cls(self): + from rbac.models import SystemRole, OrgRole + if self.scope == 'org': + return OrgRole + else: + return SystemRole @property def display(self): @@ -181,15 +197,13 @@ class RoleManager(models.Manager): @property def role_bindings(self): - from rbac.models import RoleBinding - queryset = RoleBinding.objects.filter(user=self.user) + queryset = self.role_binding_cls.objects.filter(user=self.user) if self.scope: queryset = queryset.filter(scope=self.scope) return queryset def _get_queryset(self): - from rbac.models import RoleBinding - queryset = RoleBinding.get_user_roles(self.user) + queryset = self.role_binding_cls.get_user_roles(self.user) if self.scope: queryset = queryset.filter(scope=self.scope) return queryset @@ -204,36 +218,84 @@ class RoleManager(models.Manager): return return self.role_bindings.delete() - def add(self, *roles): - from rbac.models import RoleBinding - items = [] + def _clean_roles(self, roles_or_ids): + if not roles_or_ids: + return + is_model = isinstance(roles_or_ids[0], models.Model) + if not is_model: + roles = self.role_cls.objects.filter(id__in=roles_or_ids) + else: + roles = roles_or_ids + roles = list([r for r in roles if r.scope == self.scope]) + return roles - for role in roles: + def add(self, *roles): + if not roles: + return + + roles = self._clean_roles(roles) + old_ids = self.role_bindings.values_list('role', flat=True) + need_adds = [r for r in roles if r.id not in old_ids] + + items = [] + for role in need_adds: kwargs = { 'role': role, 'user': self.user, - 'scope': role.scope + 'scope': self.scope } - if self.scope and role.scope != self.scope: - continue - if not current_org.is_root() and role.scope == RoleBinding.Scope.org: + if not current_org.is_root(): kwargs['org_id'] = current_org.id - items.append(RoleBinding(**kwargs)) + items.append(self.role_binding_cls(**kwargs)) try: - RoleBinding.objects.bulk_create(items, ignore_conflicts=True) + self.role_binding_cls.objects.bulk_create(items, ignore_conflicts=True) except Exception as e: logger.error('Create role binding error: {}'.format(e)) - def set(self, roles): - self.clear() - self.add(*roles) + def set(self, roles, clear=False): + if clear: + self.clear() + self.add(*roles) + return + + role_ids = set([r.id for r in roles]) + old_ids = self.role_bindings.values_list('role', flat=True) + old_ids = set(old_ids) + + del_ids = old_ids - role_ids + add_ids = role_ids - old_ids + self.remove(*del_ids) + self.add(*add_ids) + + def remove(self, *roles): + if not roles: + return + roles = self._clean_roles(roles) + return self.role_bindings.filter(role__in=roles).delete() def cache_set(self, roles): query = self._get_queryset() query._result_cache = roles self._cache = query + def remove_role_system_admin(self): + role = self.builtin_role.system_admin.get_role() + return self.remove(role) + + def add_role_system_admin(self): + role = self.builtin_role.system_admin.get_role() + return self.add(role) + + def add_role_system_user(self): + role = self.builtin_role.system_user.get_role() + return self.add(role) + + @property + def builtin_role(self): + from rbac.builtin import BuiltinRole + return BuiltinRole + class OrgRoleManager(RoleManager): def __init__(self, *args, **kwargs): @@ -257,6 +319,7 @@ class RoleMixin: _org_roles = None _system_roles = None PERM_CACHE_KEY = 'USER_PERMS_{}_{}' + _is_superuser = None @lazyproperty def roles(self): @@ -288,17 +351,27 @@ class RoleMixin: key = cls.PERM_CACHE_KEY.format('*', '*') cache.delete_pattern(key) - @lazyproperty + @property def is_superuser(self): """ 由于这里用了 cache ，所以不能改成 self.system_roles.filter().exists() 会查询的 """ + if not self._is_superuser: + return self._is_superuser + from rbac.builtin import BuiltinRole - # return self.system_roles.all().filter(id=BuiltinRole.system_admin.id).exists() ids = [str(r.id) for r in self.system_roles.all()] yes = BuiltinRole.system_admin.id in ids + self._is_superuser = yes return yes + @is_superuser.setter + def is_superuser(self, value): + if value: + self.system_roles.add_role_system_admin() + else: + self.system_roles.remove_role_system_admin() + @lazyproperty def is_org_admin(self): from rbac.builtin import BuiltinRole @@ -382,11 +455,6 @@ class RoleMixin: perms = RoleBinding.get_user_perms(self) return perms - def set_default_system_role(self): - from rbac.builtin import BuiltinRole - role_user = BuiltinRole.system_user.get_role() - self.system_roles.add(role_user) - class TokenMixin: CACHE_KEY_USER_RESET_PASSWORD_PREFIX = "_KEY_USER_RESET_PASSWORD_{}" diff --git a/apps/users/signal_handlers.py b/apps/users/signal_handlers.py index 8f2065ca3..065386fde 100644 --- a/apps/users/signal_handlers.py +++ b/apps/users/signal_handlers.py @@ -68,7 +68,7 @@ def on_user_create_set_default_system_role(sender, instance, created, **kwargs): has_system_role = instance.system_roles.all().exists() if not has_system_role: logger.debug("Receive user create signal, set default role") - instance.set_default_system_role() + instance.system_roles.add_role_system_user() @receiver(post_user_create)