github
/
jumpserver
mirror of https://github.com/jumpserver/jumpserver
1
0
Fork 0
Code Issues Releases Wiki Activity

perf: 优化 role bingding，优化 is_superuser

pull/7950/head
ibuler 2022-03-21 16:40:14 +08:00 committed by Jiangjie.Bai
parent 3345456dc2
commit 87c6eec619
4 changed files with 110 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
apps/rbac/models/role.py
View File

@ -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()

14
apps/rbac/models/rolebinding.py
View File

@ -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)

116
apps/users/models/user.py
View File

@ -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_{}"

2
apps/users/signal_handlers.py
View File

@ -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)