From d3be16ffe8d6e354ca4ca6ac1a77c42e8369429f Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Sun, 27 Sep 2020 16:02:44 +0800 Subject: [PATCH] fix (#4680) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(perms): 资产授权列表关联数据改为 `prefetch_related` * perf(perms): 优化一波 * dispatch_mapping_node_tasks.delay * perf: 在做一些优化 * perf: 再优化一波 * perf(perms): 授权更改节点慢的问题 * fix: 修改一处bug * perf(perms): ungrouped 资产数量计算方式 * fix: 修复dispatch data中的bug * fix(assets): add_nodes_assets_to_system_users celery task * fix: 修复ungrouped的bug * feat(nodes): 添加 favorite 节点 * feat(node): 添加 favorite api * fix: 修复clean keys的bug Co-authored-by: xinwen Co-authored-by: ibuler --- apps/assets/api/mixin.py | 2 +- apps/assets/models/asset.py | 5 + apps/assets/models/favorite_asset.py | 8 + apps/assets/models/node.py | 146 +++++------ apps/assets/pagination.py | 9 +- apps/assets/serializers/asset.py | 3 +- apps/assets/tasks/__init__.py | 1 + apps/assets/tasks/nodes_amount.py | 14 ++ apps/assets/utils.py | 9 +- apps/jumpserver/conf.py | 1 + apps/jumpserver/settings/custom.py | 1 + apps/ops/celery/__init__.py | 14 ++ apps/perms/api/asset_permission_relation.py | 2 +- apps/perms/api/user_permission/mixin.py | 43 ++-- .../user_permission/user_permission_assets.py | 114 ++++----- .../user_permission/user_permission_nodes.py | 120 +++++---- .../user_permission_nodes_with_assets.py | 179 ++++++-------- apps/perms/models/asset_permission.py | 27 +- apps/perms/models/base.py | 6 +- apps/perms/pagination.py | 16 +- apps/perms/serializers/asset_permission.py | 6 +- apps/perms/signals_handler.py | 25 +- apps/perms/tasks.py | 40 ++- apps/perms/urls/asset_permission.py | 12 +- apps/perms/utils/__init__.py | 2 +- apps/perms/utils/asset_permission.py | 37 +-- ..._node_tree.py => user_asset_permission.py} | 233 +++++++++++++++--- jms | 9 +- 28 files changed, 627 insertions(+), 457 deletions(-) rename apps/perms/utils/{user_node_tree.py => user_asset_permission.py} (58%) diff --git a/apps/assets/api/mixin.py b/apps/assets/api/mixin.py index 2876fd0e3..01dd42805 100644 --- a/apps/assets/api/mixin.py +++ b/apps/assets/api/mixin.py @@ -85,7 +85,7 @@ class FilterAssetByNodeMixin: show_current_asset_arg = request.query_params.get('show_current_asset') if show_current_asset_arg is not None: return show_current_asset_arg != '1' - return query_all_arg == '1' + return query_all_arg != '0' @lazyproperty def node(self): diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index 82d09944c..65414e86d 100644 --- a/apps/assets/models/asset.py +++ b/apps/assets/models/asset.py @@ -47,6 +47,10 @@ class AssetManager(OrgManager): ) +class AssetOrgManager(OrgManager): + pass + + class AssetQuerySet(models.QuerySet): def active(self): return self.filter(is_active=True) @@ -226,6 +230,7 @@ class Asset(ProtocolsMixin, NodesRelationMixin, OrgModelMixin): comment = models.TextField(max_length=128, default='', blank=True, verbose_name=_('Comment')) objects = AssetManager.from_queryset(AssetQuerySet)() + org_objects = AssetOrgManager.from_queryset(AssetQuerySet)() _connectivity = None def __str__(self): diff --git a/apps/assets/models/favorite_asset.py b/apps/assets/models/favorite_asset.py index 3abc69c8c..af47a867f 100644 --- a/apps/assets/models/favorite_asset.py +++ b/apps/assets/models/favorite_asset.py @@ -18,3 +18,11 @@ class FavoriteAsset(CommonModelMixin): @classmethod def get_user_favorite_assets_id(cls, user): return cls.objects.filter(user=user).values_list('asset', flat=True) + + @classmethod + def get_user_favorite_assets(cls, user): + from assets.models import Asset + from perms.utils.user_asset_permission import get_user_granted_all_assets + asset_ids = get_user_granted_all_assets(user).values_list('id', flat=True) + query_name = cls.asset.field.related_query_name() + return Asset.org_objects.filter(**{f'{query_name}__user_id': user.id}, id__in=asset_ids).distinct() diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index dd91cd993..5b4dd7259 100644 --- a/apps/assets/models/node.py +++ b/apps/assets/models/node.py @@ -41,16 +41,16 @@ class FamilyMixin: @staticmethod def clean_children_keys(nodes_keys): - nodes_keys = sorted(list(nodes_keys), key=lambda x: (len(x), x)) + sort_key = lambda k: [int(i) for i in k.split(':')] + nodes_keys = sorted(list(nodes_keys), key=sort_key) + nodes_keys_clean = [] - for key in nodes_keys[::-1]: - found = False - for k in nodes_keys: - if key.startswith(k + ':'): - found = True - break - if not found: - nodes_keys_clean.append(key) + base_key = '' + for key in nodes_keys: + if key.startswith(base_key + ':'): + continue + nodes_keys_clean.append(key) + base_key = key return nodes_keys_clean @classmethod @@ -213,26 +213,29 @@ class NodeAssetsMixin: key = '' id = None - @classmethod - def clear_all_nodes_assets_amount(cls): - nodes = cls.objects.all() - for node in nodes: - count = node.get_all_assets().count() - def get_all_assets(self): from .asset import Asset - if self.is_org_root(): - return Asset.objects.filter(org_id=self.org_id) - - q = Q(nodes__key__startswith=self.key) | Q(nodes__key=self.key) + q = Q(nodes__key__startswith=f'{self.key}:') | Q(nodes__key=self.key) return Asset.objects.filter(q).distinct() + @classmethod + def get_node_all_assets_by_key_v2(cls, key): + # 最初的写法是： + # Asset.objects.filter(Q(nodes__key__startswith=f'{node.key}:') | Q(nodes__id=node.id)) + # 可是 startswith 会导致表关联时 Asset 索引失效 + from .asset import Asset + node_ids = cls.objects.filter( + Q(key__startswith=f'{key}:') | + Q(key=key) + ).values_list('id', flat=True).distinct() + assets = Asset.objects.filter( + nodes__id__in=list(node_ids) + ).distinct() + return assets + def get_assets(self): from .asset import Asset - if self.is_org_root(): - assets = Asset.objects.filter(Q(nodes=self) | Q(nodes__isnull=True)) - else: - assets = Asset.objects.filter(nodes=self) + assets = Asset.objects.filter(nodes=self) return assets.distinct() def get_valid_assets(self): @@ -241,51 +244,54 @@ class NodeAssetsMixin: def get_all_valid_assets(self): return self.get_all_assets().valid() - @classmethod - def _get_nodes_all_assets(cls, nodes_keys): - """ - 当节点比较多的时候，这种正则方式性能差极了 - :param nodes_keys: - :return: - """ - from .asset import Asset - nodes_keys = cls.clean_children_keys(nodes_keys) - nodes_children_pattern = set() - for key in nodes_keys: - children_pattern = cls.get_node_all_children_key_pattern(key) - nodes_children_pattern.add(children_pattern) - pattern = '|'.join(nodes_children_pattern) - return Asset.objects.filter(nodes__key__regex=pattern).distinct() - @classmethod def get_nodes_all_assets_ids(cls, nodes_keys): - nodes_keys = cls.clean_children_keys(nodes_keys) - assets_ids = set() - for key in nodes_keys: - node_assets_ids = cls.tree().all_assets(key) - assets_ids.update(set(node_assets_ids)) + assets_ids = cls.get_nodes_all_assets(nodes_keys).values_list('id', flat=True) return assets_ids @classmethod def get_nodes_all_assets(cls, nodes_keys, extra_assets_ids=None): from .asset import Asset nodes_keys = cls.clean_children_keys(nodes_keys) - assets_ids = cls.get_nodes_all_assets_ids(nodes_keys) + q = Q() + node_ids = () + for key in nodes_keys: + q |= Q(key__startswith=f'{key}:') + q |= Q(key=key) + if q: + node_ids = Node.objects.filter(q).distinct().values_list('id', flat=True) + + q = Q(nodes__id__in=list(node_ids)) if extra_assets_ids: - assets_ids.update(set(extra_assets_ids)) - return Asset.objects.filter(id__in=assets_ids) + q |= Q(id__in=extra_assets_ids) + if q: + return Asset.org_objects.filter(q).distinct() + else: + return Asset.objects.none() class SomeNodesMixin: key = '' default_key = '1' default_value = 'Default' - ungrouped_key = '-10' - ungrouped_value = _('ungrouped') empty_key = '-11' empty_value = _("empty") - favorite_key = '-12' - favorite_value = _("favorite") + + @classmethod + def default_node(cls): + with tmp_to_org(Organization.default()): + defaults = {'value': cls.default_value} + try: + obj, created = cls.objects.get_or_create( + defaults=defaults, key=cls.default_key, + ) + except IntegrityError as e: + logger.error("Create default node failed: {}".format(e)) + cls.modify_other_org_root_node_key() + obj, created = cls.objects.get_or_create( + defaults=defaults, key=cls.default_key, + ) + return obj def is_default_node(self): return self.key == self.default_key @@ -320,51 +326,15 @@ class SomeNodesMixin: @classmethod def org_root(cls): - root = cls.objects.filter(key__regex=r'^[0-9]+$') + root = cls.objects.filter(parent_key='').exclude(key__startswith='-') if root: return root[0] else: return cls.create_org_root_node() - @classmethod - def ungrouped_node(cls): - with tmp_to_org(Organization.system()): - defaults = {'value': cls.ungrouped_value} - obj, created = cls.objects.get_or_create( - defaults=defaults, key=cls.ungrouped_key - ) - return obj - - @classmethod - def default_node(cls): - with tmp_to_org(Organization.default()): - defaults = {'value': cls.default_value} - try: - obj, created = cls.objects.get_or_create( - defaults=defaults, key=cls.default_key, - ) - except IntegrityError as e: - logger.error("Create default node failed: {}".format(e)) - cls.modify_other_org_root_node_key() - obj, created = cls.objects.get_or_create( - defaults=defaults, key=cls.default_key, - ) - return obj - - @classmethod - def favorite_node(cls): - with tmp_to_org(Organization.system()): - defaults = {'value': cls.favorite_value} - obj, created = cls.objects.get_or_create( - defaults=defaults, key=cls.favorite_key - ) - return obj - @classmethod def initial_some_nodes(cls): cls.default_node() - cls.ungrouped_node() - cls.favorite_node() @classmethod def modify_other_org_root_node_key(cls): diff --git a/apps/assets/pagination.py b/apps/assets/pagination.py index 0a2562c29..dd37600f8 100644 --- a/apps/assets/pagination.py +++ b/apps/assets/pagination.py @@ -17,16 +17,13 @@ class AssetLimitOffsetPagination(LimitOffsetPagination): exclude_query_params = { self.limit_query_param, self.offset_query_param, - 'node', 'all', 'show_current_asset' + 'node', 'all', 'show_current_asset', + 'node_id', 'display', 'draw', } - has_filter = False for k, v in self._request.query_params.items(): if k not in exclude_query_params and v is not None: - has_filter = True - break - if has_filter: - return super().get_count(queryset) + return super().get_count(queryset) is_query_all = self._view.is_query_node_all_assets if is_query_all: diff --git a/apps/assets/serializers/asset.py b/apps/assets/serializers/asset.py index 1f03d14c9..75fc220e0 100644 --- a/apps/assets/serializers/asset.py +++ b/apps/assets/serializers/asset.py @@ -115,6 +115,7 @@ class AssetSerializer(BulkOrgResourceModelSerializer): def setup_eager_loading(cls, queryset): """ Perform necessary eager loading of data. """ queryset = queryset.select_related('admin_user', 'domain', 'platform') + queryset = queryset.prefetch_related('nodes', 'labels') return queryset def compatible_with_old_protocol(self, validated_data): @@ -152,7 +153,7 @@ class AssetDisplaySerializer(AssetSerializer): @classmethod def setup_eager_loading(cls, queryset): - """ Perform necessary eager loading of data. """ + queryset = super().setup_eager_loading(queryset) queryset = queryset\ .annotate(admin_user_username=F('admin_user__username')) return queryset diff --git a/apps/assets/tasks/__init__.py b/apps/assets/tasks/__init__.py index 6f53c9fa2..b1866d5ec 100644 --- a/apps/assets/tasks/__init__.py +++ b/apps/assets/tasks/__init__.py @@ -9,3 +9,4 @@ from .gather_asset_users import * from .gather_asset_hardware_info import * from .push_system_user import * from .system_user_connectivity import * +from .nodes_amount import * diff --git a/apps/assets/tasks/nodes_amount.py b/apps/assets/tasks/nodes_amount.py index e69de29bb..4d53be525 100644 --- a/apps/assets/tasks/nodes_amount.py +++ b/apps/assets/tasks/nodes_amount.py @@ -0,0 +1,14 @@ +from celery import shared_task + +from assets.utils import check_node_assets_amount +from common.utils import get_logger +from common.utils.timezone import now + +logger = get_logger(__file__) + + +@shared_task() +def check_node_assets_amount_celery_task(): + logger.info(f'>>> {now()} begin check_node_assets_amount_celery_task ...') + check_node_assets_amount() + logger.info(f'>>> {now()} end check_node_assets_amount_celery_task ...') diff --git a/apps/assets/utils.py b/apps/assets/utils.py index ac514ba49..c67cc8528 100644 --- a/apps/assets/utils.py +++ b/apps/assets/utils.py @@ -1,13 +1,8 @@ # ~*~ coding: utf-8 ~*~ # -from treelib import Tree -from treelib.exceptions import NodeIDAbsentError -from collections import defaultdict -from copy import deepcopy - from django.db.models import Q -from common.utils import get_logger, timeit, lazyproperty +from common.utils import get_logger from .models import Asset, Node @@ -21,7 +16,7 @@ def check_node_assets_amount(): ).distinct().count() if node.assets_amount != assets_amount: - print(f' wrong assets amount ' + print(f'>>> wrong assets amount ' f'{node.assets_amount} right is {assets_amount}') node.assets_amount = assets_amount node.save() diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py index 68c23301d..3e0b65172 100644 --- a/apps/jumpserver/conf.py +++ b/apps/jumpserver/conf.py @@ -257,6 +257,7 @@ class Config(dict): 'SYSLOG_FACILITY': 'user', 'SYSLOG_SOCKTYPE': 2, 'PERM_SINGLE_ASSET_TO_UNGROUP_NODE': False, + 'PERM_EXPIRED_CHECK_PERIODIC': 60, 'WINDOWS_SSH_DEFAULT_SHELL': 'cmd', 'FLOWER_URL': "127.0.0.1:5555", 'DEFAULT_ORG_SHOW_ALL_USERS': True, diff --git a/apps/jumpserver/settings/custom.py b/apps/jumpserver/settings/custom.py index 97479115f..47bf4bebe 100644 --- a/apps/jumpserver/settings/custom.py +++ b/apps/jumpserver/settings/custom.py @@ -75,6 +75,7 @@ BACKEND_ASSET_USER_AUTH_VAULT = False DEFAULT_ORG_SHOW_ALL_USERS = CONFIG.DEFAULT_ORG_SHOW_ALL_USERS PERM_SINGLE_ASSET_TO_UNGROUP_NODE = CONFIG.PERM_SINGLE_ASSET_TO_UNGROUP_NODE +PERM_EXPIRED_CHECK_PERIODIC = CONFIG.PERM_EXPIRED_CHECK_PERIODIC WINDOWS_SSH_DEFAULT_SHELL = CONFIG.WINDOWS_SSH_DEFAULT_SHELL FLOWER_URL = CONFIG.FLOWER_URL diff --git a/apps/ops/celery/__init__.py b/apps/ops/celery/__init__.py index e0c97816a..b8ed56be1 100644 --- a/apps/ops/celery/__init__.py +++ b/apps/ops/celery/__init__.py @@ -4,6 +4,7 @@ import os from kombu import Exchange, Queue from celery import Celery +from celery.schedules import crontab # set the default Django settings module for the 'celery' program. os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'jumpserver.settings') @@ -28,3 +29,16 @@ configs["CELERY_ROUTES"] = { app.namespace = 'CELERY' app.conf.update(configs) app.autodiscover_tasks(lambda: [app_config.split('.')[0] for app_config in settings.INSTALLED_APPS]) + +app.conf.beat_schedule = { + 'check-asset-permission-expired': { + 'task': 'perms.tasks.check_asset_permission_expired', + 'schedule': settings.PERM_EXPIRED_CHECK_PERIODIC, + 'args': () + }, + 'check-node-assets-amount': { + 'task': 'assets.tasks.nodes_amount.check_node_assets_amount_celery_task', + 'schedule': crontab(minute=0, hour=0), + 'args': () + }, +} diff --git a/apps/perms/api/asset_permission_relation.py b/apps/perms/api/asset_permission_relation.py index 172391dc8..e9309c7e7 100644 --- a/apps/perms/api/asset_permission_relation.py +++ b/apps/perms/api/asset_permission_relation.py @@ -111,7 +111,7 @@ class AssetPermissionAllAssetListApi(generics.ListAPIView): asset_q |= Q(nodes__key__startswith=f'{key}:') asset_q |= Q(nodes__key=key) - assets = Asset.objects.filter(asset_q).only(*self.serializer_class.Meta.only_fields) + assets = Asset.objects.filter(asset_q).only(*self.serializer_class.Meta.only_fields).distinct() return assets diff --git a/apps/perms/api/user_permission/mixin.py b/apps/perms/api/user_permission/mixin.py index 9466aa78c..6caffaf33 100644 --- a/apps/perms/api/user_permission/mixin.py +++ b/apps/perms/api/user_permission/mixin.py @@ -1,45 +1,41 @@ # -*- coding: utf-8 -*- # +from rest_framework.request import Request + from common.permissions import IsOrgAdminOrAppUser, IsValidUser from common.utils import lazyproperty -from rest_framework.generics import get_object_or_404 - from users.models import User from perms.models import UserGrantedMappingNode -from common.exceptions import JMSObjectDoesNotExist -from perms.async_tasks.mapping_node_task import submit_update_mapping_node_task_for_user -from ...hands import Node -class UserGrantedNodeDispatchMixin: +class UserNodeGrantStatusDispatchMixin: - def submit_update_mapping_node_task(self, user): - submit_update_mapping_node_task_for_user(user) + @staticmethod + def get_mapping_node_by_key(key): + return UserGrantedMappingNode.objects.get(key=key) - def dispatch_node_process(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): - if mapping_node is None: - ancestor_keys = Node.get_node_ancestor_keys(key) - granted = UserGrantedMappingNode.objects.filter(key__in=ancestor_keys, granted=True).exists() - if not granted: - raise JMSObjectDoesNotExist(object_name=Node._meta.object_name) - queryset = self.on_granted_node(key, mapping_node, node) + def dispatch_get_data(self, key, user): + status = UserGrantedMappingNode.get_node_granted_status(key, user) + if status == UserGrantedMappingNode.GRANTED_DIRECT: + return self.get_data_on_node_direct_granted(key) + elif status == UserGrantedMappingNode.GRANTED_INDIRECT: + return self.get_data_on_node_indirect_granted(key) else: - if mapping_node.granted: - # granted_node - queryset = self.on_granted_node(key, mapping_node, node) - else: - queryset = self.on_ungranted_node(key, mapping_node, node) - return queryset + return self.get_data_on_node_not_granted(key) - def on_granted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): + def get_data_on_node_direct_granted(self, key): raise NotImplementedError - def on_ungranted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): + def get_data_on_node_indirect_granted(self, key): + raise NotImplementedError + + def get_data_on_node_not_granted(self, key): raise NotImplementedError class ForAdminMixin: permission_classes = (IsOrgAdminOrAppUser,) + kwargs: dict @lazyproperty def user(self): @@ -49,6 +45,7 @@ class ForAdminMixin: class ForUserMixin: permission_classes = (IsValidUser,) + request: Request @lazyproperty def user(self): diff --git a/apps/perms/api/user_permission/user_permission_assets.py b/apps/perms/api/user_permission/user_permission_assets.py index cba95dd53..3cb211434 100644 --- a/apps/perms/api/user_permission/user_permission_assets.py +++ b/apps/perms/api/user_permission/user_permission_assets.py @@ -1,39 +1,32 @@ # -*- coding: utf-8 -*- # -from django.db.models import Q from django.utils.decorators import method_decorator -from perms.api.user_permission.mixin import UserGrantedNodeDispatchMixin +from perms.api.user_permission.mixin import UserNodeGrantStatusDispatchMixin from rest_framework.generics import ListAPIView from rest_framework.response import Response from django.conf import settings from assets.api.mixin import SerializeToTreeNodeMixin -from common.utils import get_object_or_none -from users.models import User -from common.permissions import IsOrgAdminOrAppUser, IsValidUser from common.utils import get_logger -from ...hands import Node -from ... import serializers -from perms.models import UserGrantedMappingNode -from perms.utils.user_node_tree import get_node_all_granted_assets from perms.pagination import GrantedAssetLimitOffsetPagination -from assets.models import Asset +from assets.models import Asset, Node, FavoriteAsset from orgs.utils import tmp_to_root_org +from ... import serializers +from ...utils.user_asset_permission import ( + get_node_all_granted_assets, get_user_direct_granted_assets, + get_user_granted_all_assets +) from .mixin import ForAdminMixin, ForUserMixin logger = get_logger(__name__) -__all__ = [ - 'UserDirectGrantedAssetsForAdminApi', 'MyAllAssetsAsTreeApi', - 'UserGrantedNodeAssetsForAdminApi', 'MyDirectGrantedAssetsApi', - 'UserDirectGrantedAssetsAsTreeForAdminApi', 'MyGrantedNodeAssetsApi', - 'MyUngroupAssetsAsTreeApi', -] - @method_decorator(tmp_to_root_org(), name='list') class UserDirectGrantedAssetsApi(ListAPIView): + """ + 用户直接授权的资产的列表，也就是授权规则上直接授权的资产，并非是来自节点的 + """ serializer_class = serializers.AssetGrantedSerializer only_fields = serializers.AssetGrantedSerializer.Meta.only_fields filter_fields = ['hostname', 'ip', 'id', 'comment'] @@ -41,17 +34,32 @@ class UserDirectGrantedAssetsApi(ListAPIView): def get_queryset(self): user = self.user + assets = get_user_direct_granted_assets(user)\ + .prefetch_related('platform')\ + .only(*self.only_fields) + return assets - return Asset.objects.filter( - Q(granted_by_permissions__users=user) | - Q(granted_by_permissions__user_groups__users=user) - ).distinct().only( - *self.only_fields - ) + +@method_decorator(tmp_to_root_org(), name='list') +class UserFavoriteGrantedAssetsApi(ListAPIView): + serializer_class = serializers.AssetGrantedSerializer + only_fields = serializers.AssetGrantedSerializer.Meta.only_fields + filter_fields = ['hostname', 'ip', 'id', 'comment'] + search_fields = ['hostname', 'ip', 'comment'] + + def get_queryset(self): + user = self.user + assets = FavoriteAsset.get_user_favorite_assets(user)\ + .prefetch_related('platform')\ + .only(*self.only_fields) + return assets @method_decorator(tmp_to_root_org(), name='list') class AssetsAsTreeMixin(SerializeToTreeNodeMixin): + """ + 将 资产 序列化成树的结构返回 + """ def list(self, request, *args, **kwargs): queryset = self.filter_queryset(self.get_queryset()) data = self.serialize_assets(queryset, None) @@ -66,6 +74,14 @@ class MyDirectGrantedAssetsApi(ForUserMixin, UserDirectGrantedAssetsApi): pass +class UserFavoriteGrantedAssetsForAdminApi(ForAdminMixin, UserFavoriteGrantedAssetsApi): + pass + + +class MyFavoriteGrantedAssetsApi(ForUserMixin, UserFavoriteGrantedAssetsApi): + pass + + @method_decorator(tmp_to_root_org(), name='list') class UserDirectGrantedAssetsAsTreeForAdminApi(ForAdminMixin, AssetsAsTreeMixin, UserDirectGrantedAssetsApi): pass @@ -85,27 +101,8 @@ class UserAllGrantedAssetsApi(ListAPIView): only_fields = serializers.AssetGrantedSerializer.Meta.only_fields def get_queryset(self): - user = self.user - - granted_node_keys = Node.objects.filter( - Q(granted_by_permissions__users=user) | - Q(granted_by_permissions__user_groups__users=user) - ).distinct().values_list('key', flat=True) - - granted_node_q = Q() - for _key in granted_node_keys: - granted_node_q |= Q(nodes__key__startswith=f'{_key}:') - granted_node_q |= Q(nodes__key=_key) - - q = Q(granted_by_permissions__users=user) | \ - Q(granted_by_permissions__user_groups__users=user) - - if granted_node_q: - q |= granted_node_q - - return Asset.objects.filter(q).distinct().only( - *self.only_fields - ) + queryset = get_user_granted_all_assets(self.user) + return queryset.only(*self.only_fields) class MyAllAssetsAsTreeApi(ForUserMixin, AssetsAsTreeMixin, UserAllGrantedAssetsApi): @@ -113,33 +110,30 @@ class MyAllAssetsAsTreeApi(ForUserMixin, AssetsAsTreeMixin, UserAllGrantedAssets @method_decorator(tmp_to_root_org(), name='list') -class UserGrantedNodeAssetsApi(UserGrantedNodeDispatchMixin, ListAPIView): +class UserGrantedNodeAssetsApi(UserNodeGrantStatusDispatchMixin, ListAPIView): serializer_class = serializers.AssetGrantedSerializer only_fields = serializers.AssetGrantedSerializer.Meta.only_fields filter_fields = ['hostname', 'ip', 'id', 'comment'] search_fields = ['hostname', 'ip', 'comment'] pagination_class = GrantedAssetLimitOffsetPagination + pagination_node: Node def get_queryset(self): node_id = self.kwargs.get("node_id") - user = self.user - - mapping_node: UserGrantedMappingNode = get_object_or_none( - UserGrantedMappingNode, user=user, node_id=node_id) node = Node.objects.get(id=node_id) - return self.dispatch_node_process(node.key, mapping_node, node) + self.pagination_node = node + return self.dispatch_get_data(node.key, self.user) - def on_granted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): - self.node = node - return Asset.objects.filter( - Q(nodes__key__startswith=f'{node.key}:') | - Q(nodes__id=node.id) - ).distinct() + def get_data_on_node_direct_granted(self, key): + # 如果这个节点是直接授权的(或者说祖先节点直接授权的), 获取下面的所有资产 + return Node.get_node_all_assets_by_key_v2(key) - def on_ungranted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): - self.node = mapping_node - user = self.user - return get_node_all_granted_assets(user, node.key) + def get_data_on_node_indirect_granted(self, key): + self.pagination_node = self.get_mapping_node_by_key(key) + return get_node_all_granted_assets(self.user, key) + + def get_data_on_node_not_granted(self, key): + return Asset.objects.none() class UserGrantedNodeAssetsForAdminApi(ForAdminMixin, UserGrantedNodeAssetsApi): diff --git a/apps/perms/api/user_permission/user_permission_nodes.py b/apps/perms/api/user_permission/user_permission_nodes.py index 39bc8dfe8..c6c98a1f2 100644 --- a/apps/perms/api/user_permission/user_permission_nodes.py +++ b/apps/perms/api/user_permission/user_permission_nodes.py @@ -1,25 +1,24 @@ # -*- coding: utf-8 -*- # -from django.db.models import Q, F -from perms.api.user_permission.mixin import ForAdminMixin, ForUserMixin +import abc from rest_framework.generics import ( ListAPIView ) from rest_framework.response import Response +from rest_framework.request import Request -from perms.utils.user_node_tree import ( - node_annotate_mapping_node, get_ungranted_node_children, - is_granted, get_granted_assets_amount, node_annotate_set_granted, -) -from common.utils.django import get_object_or_none -from common.utils import lazyproperty -from perms.models import UserGrantedMappingNode from orgs.utils import tmp_to_root_org from assets.api.mixin import SerializeToTreeNodeMixin from common.utils import get_logger -from ...hands import Node -from .mixin import UserGrantedNodeDispatchMixin +from .mixin import ForAdminMixin, ForUserMixin, UserNodeGrantStatusDispatchMixin +from ...hands import Node, User from ... import serializers +from ...utils.user_asset_permission import ( + get_indirect_granted_node_children, + get_user_granted_nodes_list_via_mapping_node, + get_top_level_granted_nodes, + init_user_tree_if_need, +) logger = get_logger(__name__) @@ -32,12 +31,13 @@ __all__ = [ 'MyGrantedNodeChildrenApi', 'UserGrantedNodeChildrenAsTreeForAdminApi', 'MyGrantedNodeChildrenAsTreeApi', - 'NodeChildrenAsTreeApi', + 'BaseGrantedNodeAsTreeApi', + 'UserGrantedNodesMixin', ] -class GrantedNodeBaseApi(ListAPIView): - @lazyproperty +class _GrantedNodeStructApi(ListAPIView, metaclass=abc.ABCMeta): + @property def user(self): raise NotImplementedError @@ -47,113 +47,105 @@ class GrantedNodeBaseApi(ListAPIView): raise NotImplementedError -class NodeChildrenApi(GrantedNodeBaseApi): +class NodeChildrenMixin: + def get_children(self): + raise NotImplementedError + + def get_nodes(self): + nodes = self.get_children() + return nodes + + +class BaseGrantedNodeApi(_GrantedNodeStructApi, metaclass=abc.ABCMeta): serializer_class = serializers.NodeGrantedSerializer @tmp_to_root_org() def list(self, request, *args, **kwargs): + init_user_tree_if_need(self.user) nodes = self.get_nodes() serializer = self.get_serializer(nodes, many=True) return Response(serializer.data) -class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, GrantedNodeBaseApi): +class BaseNodeChildrenApi(NodeChildrenMixin, BaseGrantedNodeApi, metaclass=abc.ABCMeta): + pass + + +class BaseGrantedNodeAsTreeApi(SerializeToTreeNodeMixin, _GrantedNodeStructApi, metaclass=abc.ABCMeta): @tmp_to_root_org() def list(self, request, *args, **kwargs): + init_user_tree_if_need(self.user) nodes = self.get_nodes() nodes = self.serialize_nodes(nodes, with_asset_amount=True) return Response(data=nodes) -class UserGrantedNodeChildrenMixin(UserGrantedNodeDispatchMixin): +class BaseNodeChildrenAsTreeApi(NodeChildrenMixin, BaseGrantedNodeAsTreeApi, metaclass=abc.ABCMeta): + pass - def get_nodes(self): + +class UserGrantedNodeChildrenMixin(UserNodeGrantStatusDispatchMixin): + user: User + request: Request + + def get_children(self): user = self.user key = self.request.query_params.get('key') - self.submit_update_mapping_node_task(user) - if not key: - nodes = get_ungranted_node_children(user) + nodes = list(get_top_level_granted_nodes(user)) else: - mapping_node = get_object_or_none( - UserGrantedMappingNode, user=user, key=key - ) - nodes = self.dispatch_node_process(key, mapping_node, None) + nodes = self.dispatch_get_data(key, user) return nodes - def on_granted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): + def get_data_on_node_direct_granted(self, key): return Node.objects.filter(parent_key=key) - def on_ungranted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): - user = self.user - nodes = get_ungranted_node_children(user, key) + def get_data_on_node_indirect_granted(self, key): + nodes = get_indirect_granted_node_children(self.user, key) return nodes + def get_data_on_node_not_granted(self, key): + return Node.objects.none() + class UserGrantedNodesMixin: """ 查询用户授权的所有节点 直接授权节点 + 授权资产关联的节点 """ + user: User def get_nodes(self): - user = self.user - - # 获取 `UserGrantedMappingNode` 中对应的 `Node` - nodes = Node.objects.filter( - mapping_nodes__user=user, - ).annotate(**node_annotate_mapping_node).distinct() - - key2nodes_mapper = {} - descendant_q = Q() - - for _node in nodes: - if not is_granted(_node): - # 未授权的节点资产数量设置为 `UserGrantedMappingNode` 中的数量 - _node.assets_amount = get_granted_assets_amount(_node) - else: - # 直接授权的节点 - # 增加查询后代节点的过滤条件 - descendant_q |= Q(key__startswith=f'{_node.key}:') - - key2nodes_mapper[_node.key] = _node - - if descendant_q: - descendant_nodes = Node.objects.filter(descendant_q).annotate(**node_annotate_set_granted) - for _node in descendant_nodes: - key2nodes_mapper[_node.key] = _node - - all_nodes = key2nodes_mapper.values() - return all_nodes + return get_user_granted_nodes_list_via_mapping_node(self.user) # ------------------------------------------ # 最终的 api -class UserGrantedNodeChildrenForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, NodeChildrenApi): +class UserGrantedNodeChildrenForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi): pass -class MyGrantedNodeChildrenApi(ForUserMixin, UserGrantedNodeChildrenMixin, NodeChildrenApi): +class MyGrantedNodeChildrenApi(ForUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenApi): pass -class UserGrantedNodeChildrenAsTreeForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, NodeChildrenAsTreeApi): +class UserGrantedNodeChildrenAsTreeForAdminApi(ForAdminMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi): pass -class MyGrantedNodeChildrenAsTreeApi(ForUserMixin, UserGrantedNodeChildrenMixin, NodeChildrenAsTreeApi): +class MyGrantedNodeChildrenAsTreeApi(ForUserMixin, UserGrantedNodeChildrenMixin, BaseNodeChildrenAsTreeApi): pass -class UserGrantedNodesForAdminApi(ForAdminMixin, UserGrantedNodesMixin, NodeChildrenApi): +class UserGrantedNodesForAdminApi(ForAdminMixin, UserGrantedNodesMixin, BaseGrantedNodeApi): pass -class MyGrantedNodesApi(ForUserMixin, UserGrantedNodesMixin, NodeChildrenApi): +class MyGrantedNodesApi(ForUserMixin, UserGrantedNodesMixin, BaseGrantedNodeApi): pass -class MyGrantedNodesAsTreeApi(ForUserMixin, UserGrantedNodesMixin, NodeChildrenAsTreeApi): +class MyGrantedNodesAsTreeApi(ForUserMixin, UserGrantedNodesMixin, BaseGrantedNodeAsTreeApi): pass # ------------------------------------------ diff --git a/apps/perms/api/user_permission/user_permission_nodes_with_assets.py b/apps/perms/api/user_permission/user_permission_nodes_with_assets.py index 88f0202c5..09cfbe75f 100644 --- a/apps/perms/api/user_permission/user_permission_nodes_with_assets.py +++ b/apps/perms/api/user_permission/user_permission_nodes_with_assets.py @@ -3,80 +3,44 @@ from rest_framework.generics import ListAPIView from rest_framework.request import Request from rest_framework.response import Response -from django.db.models import Q, F -from users.models import User -from common.permissions import IsValidUser, IsOrgAdminOrAppUser -from common.utils.django import get_object_or_none -from common.utils import get_logger -from .user_permission_nodes import MyGrantedNodesAsTreeApi -from .mixin import UserGrantedNodeDispatchMixin -from perms.models import UserGrantedMappingNode -from perms.utils.user_node_tree import ( - TMP_GRANTED_FIELD, TMP_GRANTED_ASSETS_AMOUNT_FIELD, node_annotate_mapping_node, - is_granted, get_granted_assets_amount, node_annotate_set_granted, - get_granted_q, get_ungranted_node_children +from common.permissions import IsValidUser +from common.utils import get_logger, get_object_or_none +from .mixin import UserNodeGrantStatusDispatchMixin, ForUserMixin, ForAdminMixin +from ...utils.user_asset_permission import ( + get_user_resources_q_granted_by_permissions, + get_indirect_granted_node_children, UNGROUPED_NODE_KEY, FAVORITE_NODE_KEY, + get_user_direct_granted_assets, get_top_level_granted_nodes, + get_user_granted_nodes_list_via_mapping_node, + get_user_granted_all_assets, init_user_tree_if_need, + get_user_all_assetpermission_ids, ) -from assets.models import Asset +from assets.models import Asset, FavoriteAsset from assets.api import SerializeToTreeNodeMixin from orgs.utils import tmp_to_root_org from ...hands import Node logger = get_logger(__name__) -__all__ = [ - 'MyGrantedNodesAsTreeApi', - 'UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi', - 'MyGrantedNodesWithAssetsAsTreeApi', - 'MyGrantedNodeChildrenWithAssetsAsTreeApi', -] - class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView): permission_classes = (IsValidUser,) @tmp_to_root_org() def list(self, request: Request, *args, **kwargs): + """ + 此算法依赖 UserGrantedMappingNode + 获取所有授权的节点和资产 + + Node = UserGrantedMappingNode + 授权节点的子节点 + Asset = 授权节点的资产 + 直接授权的资产 + """ + user = request.user - - # 获取 `UserGrantedMappingNode` 中对应的 `Node` - nodes = Node.objects.filter( - mapping_nodes__user=user, - ).annotate(**node_annotate_mapping_node).distinct() - - key2nodes_mapper = {} - descendant_q = Q() - granted_q = Q() - - for _node in nodes: - if not is_granted(_node): - # 未授权的节点资产数量设置为 `UserGrantedMappingNode` 中的数量 - _node.assets_amount = get_granted_assets_amount(_node) - else: - # 直接授权的节点 - - # 增加查询该节点及其后代节点资产的过滤条件 - granted_q |= Q(nodes__key__startswith=f'{_node.key}:') - granted_q |= Q(nodes__key=_node.key) - - # 增加查询后代节点的过滤条件 - descendant_q |= Q(key__startswith=f'{_node.key}:') - - key2nodes_mapper[_node.key] = _node - - if descendant_q: - descendant_nodes = Node.objects.filter(descendant_q).annotate(**node_annotate_set_granted) - for _node in descendant_nodes: - key2nodes_mapper[_node.key] = _node - - all_nodes = key2nodes_mapper.values() - - # 查询出所有资产 - all_assets = Asset.objects.filter( - get_granted_q(user) | - granted_q - ).annotate(parent_key=F('nodes__key')).distinct() + init_user_tree_if_need(user) + all_nodes = get_user_granted_nodes_list_via_mapping_node(user) + all_assets = get_user_granted_all_assets(user) data = [ *self.serialize_nodes(all_nodes, with_asset_amount=True), @@ -85,61 +49,70 @@ class MyGrantedNodesWithAssetsAsTreeApi(SerializeToTreeNodeMixin, ListAPIView): return Response(data=data) -class UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi(UserGrantedNodeDispatchMixin, SerializeToTreeNodeMixin, ListAPIView): - permission_classes = (IsOrgAdminOrAppUser, ) +class UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi(ForAdminMixin, UserNodeGrantStatusDispatchMixin, + SerializeToTreeNodeMixin, ListAPIView): + """ + 带资产的授权树 + """ - def on_granted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): + def get_data_on_node_direct_granted(self, key): nodes = Node.objects.filter(parent_key=key) - assets = Asset.objects.filter(nodes__key=key).distinct() + assets = Asset.org_objects.filter(nodes__key=key).distinct() + assets = assets.prefetch_related('platform') return nodes, assets - def on_ungranted_node(self, key, mapping_node: UserGrantedMappingNode, node: Node = None): - user = self.get_user() - assets = Asset.objects.none() - nodes = Node.objects.filter( - parent_key=key, - mapping_nodes__user=user, - ).annotate( - **node_annotate_mapping_node + def get_data_on_node_indirect_granted(self, key): + user = self.user + asset_perm_ids = get_user_all_assetpermission_ids(user) + + nodes = get_indirect_granted_node_children(user, key) + + assets = Asset.org_objects.filter( + nodes__key=key, + ).filter( + granted_by_permissions__id__in=asset_perm_ids ).distinct() - - # TODO 可配置 - for _node in nodes: - if not is_granted(_node): - _node.assets_amount = get_granted_assets_amount(_node) - - if mapping_node.asset_granted: - assets = Asset.objects.filter( - nodes__key=key, - ).filter(get_granted_q(user)) + assets = assets.prefetch_related('platform') return nodes, assets - def get_user(self): - user_id = self.kwargs.get('pk') - return User.objects.get(id=user_id) + def get_data_on_node_not_granted(self, key): + return Node.objects.none(), Asset.objects.none() + + def get_data(self, key, user): + assets, nodes = [], [] + if not key: + root_nodes = get_top_level_granted_nodes(user) + nodes.extend(root_nodes) + elif key == UNGROUPED_NODE_KEY: + assets = get_user_direct_granted_assets(user) + assets = assets.prefetch_related('platform') + elif key == FAVORITE_NODE_KEY: + assets = FavoriteAsset.get_user_favorite_assets(user) + else: + nodes, assets = self.dispatch_get_data(key, user) + return nodes, assets + + def id2key_if_have(self): + id = self.request.query_params.get('id') + if id is not None: + node = get_object_or_none(Node, id=id) + if node: + return node.key @tmp_to_root_org() def list(self, request: Request, *args, **kwargs): - user = self.get_user() - key = request.query_params.get('key') - self.submit_update_mapping_node_task(user) + key = self.request.query_params.get('key') + if key is None: + key = self.id2key_if_have() - nodes = [] - assets = [] - if not key: - root_nodes = get_ungranted_node_children(user) - nodes.extend(root_nodes) - else: - mapping_node: UserGrantedMappingNode = get_object_or_none( - UserGrantedMappingNode, user=user, key=key) - nodes, assets = self.dispatch_node_process(key, mapping_node) - nodes = self.serialize_nodes(nodes, with_asset_amount=True) - assets = self.serialize_assets(assets, key) - return Response(data=[*nodes, *assets]) + user = self.user + init_user_tree_if_need(user) + nodes, assets = self.get_data(key, user) + + tree_nodes = self.serialize_nodes(nodes, with_asset_amount=True) + tree_assets = self.serialize_assets(assets, key) + return Response(data=[*tree_nodes, *tree_assets]) -class MyGrantedNodeChildrenWithAssetsAsTreeApi(UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi): - permission_classes = (IsValidUser, ) - - def get_user(self): - return self.request.user +class MyGrantedNodeChildrenWithAssetsAsTreeApi(ForUserMixin, UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi): + pass diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py index 8ef52d6b9..5611f6ad5 100644 --- a/apps/perms/models/asset_permission.py +++ b/apps/perms/models/asset_permission.py @@ -5,7 +5,6 @@ from functools import reduce from django.utils.translation import ugettext_lazy as _ from common.db import models -from common.fields.model import JsonListTextField from common.utils import lazyproperty from orgs.models import Organization from orgs.utils import get_current_org @@ -17,6 +16,8 @@ from .base import BasePermission __all__ = [ 'AssetPermission', 'Action', 'UserGrantedMappingNode', 'RebuildUserTreeTask', ] + +# 使用场景 logger = logging.getLogger(__name__) @@ -98,6 +99,14 @@ class AssetPermission(BasePermission): verbose_name = _("Asset permission") ordering = ('name',) + @lazyproperty + def users_amount(self): + return self.users.count() + + @lazyproperty + def user_groups_amount(self): + return self.user_groups.count() + @lazyproperty def assets_amount(self): return self.assets.count() @@ -186,6 +195,22 @@ class UserGrantedMappingNode(FamilyMixin, models.JMSBaseModel): parent_key = models.CharField(max_length=64, default='', verbose_name=_('Parent key'), db_index=True) # '1:1:1:1' assets_amount = models.IntegerField(default=0) + GRANTED_DIRECT = 1 + GRANTED_INDIRECT = 2 + GRANTED_NONE = 0 + + @classmethod + def get_node_granted_status(cls, key, user): + ancestor_keys = Node.get_node_ancestor_keys(key, with_self=True) + has_granted = UserGrantedMappingNode.objects.filter( + key__in=ancestor_keys, user=user + ).values_list('granted', flat=True) + if not has_granted: + return cls.GRANTED_NONE + if any(list(has_granted)): + return cls.GRANTED_DIRECT + return cls.GRANTED_INDIRECT + class RebuildUserTreeTask(models.JMSBaseModel): user = models.ForeignKey('users.User', on_delete=models.CASCADE, verbose_name=_('User')) diff --git a/apps/perms/models/base.py b/apps/perms/models/base.py index 4ad52b2ce..f5adeb838 100644 --- a/apps/perms/models/base.py +++ b/apps/perms/models/base.py @@ -13,7 +13,7 @@ from orgs.mixins.models import OrgManager __all__ = [ - 'BasePermission', + 'BasePermission', 'BasePermissionQuerySet' ] @@ -46,8 +46,8 @@ class BasePermissionManager(OrgManager): class BasePermission(OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) name = models.CharField(max_length=128, verbose_name=_('Name')) - users = models.ManyToManyField('users.User', blank=True, verbose_name=_("User")) - user_groups = models.ManyToManyField('users.UserGroup', blank=True, verbose_name=_("User group")) + users = models.ManyToManyField('users.User', blank=True, verbose_name=_("User"), related_name='%(class)ss') + user_groups = models.ManyToManyField('users.UserGroup', blank=True, verbose_name=_("User group"), related_name='%(class)ss') is_active = models.BooleanField(default=True, verbose_name=_('Active')) date_start = models.DateTimeField(default=timezone.now, db_index=True, verbose_name=_("Date start")) date_expired = models.DateTimeField(default=date_expired_default, db_index=True, verbose_name=_('Date expired')) diff --git a/apps/perms/pagination.py b/apps/perms/pagination.py index 6c794d193..75cf6c493 100644 --- a/apps/perms/pagination.py +++ b/apps/perms/pagination.py @@ -1,6 +1,10 @@ from rest_framework.pagination import LimitOffsetPagination from rest_framework.request import Request +from common.utils import get_logger + +logger = get_logger(__name__) + class GrantedAssetLimitOffsetPagination(LimitOffsetPagination): def get_count(self, queryset): @@ -10,15 +14,15 @@ class GrantedAssetLimitOffsetPagination(LimitOffsetPagination): 'key', 'all', 'show_current_asset', 'cache_policy', 'display', 'draw' } - has_filter = False for k, v in self._request.query_params.items(): if k not in exclude_query_params and v is not None: - has_filter = True - break - if has_filter: + return super().get_count(queryset) + node = getattr(self._view, 'pagination_node', None) + if node: + logger.debug(f'{self._request.get_full_path()} hit node.assets_amount[{node.assets_amount}]') + return node.assets_amount + else: return super().get_count(queryset) - node = self._view.node - return node.assets_amount def paginate_queryset(self, queryset, request: Request, view=None): self._request = request diff --git a/apps/perms/serializers/asset_permission.py b/apps/perms/serializers/asset_permission.py index a256a7a3c..2c2cfd3ce 100644 --- a/apps/perms/serializers/asset_permission.py +++ b/apps/perms/serializers/asset_permission.py @@ -56,9 +56,5 @@ class AssetPermissionSerializer(BulkOrgResourceModelSerializer): @classmethod def setup_eager_loading(cls, queryset): """ Perform necessary eager loading of data. """ - queryset = queryset.annotate( - users_amount=Count('users', distinct=True), user_groups_amount=Count('user_groups', distinct=True), - assets_amount=Count('assets', distinct=True), nodes_amount=Count('nodes', distinct=True), - system_users_amount=Count('system_users', distinct=True) - ) + queryset = queryset.prefetch_related('users', 'user_groups', 'assets', 'nodes', 'system_users') return queryset diff --git a/apps/perms/signals_handler.py b/apps/perms/signals_handler.py index 84fa02561..f2ba8e803 100644 --- a/apps/perms/signals_handler.py +++ b/apps/perms/signals_handler.py @@ -2,13 +2,13 @@ # from itertools import chain -from django.db.models.signals import m2m_changed, pre_delete +from django.db.models.signals import m2m_changed, pre_delete, pre_save from django.dispatch import receiver from django.db import transaction from django.db.models import Q from perms.tasks import dispatch_mapping_node_tasks -from users.models import User +from users.models import User, UserGroup from assets.models import Asset from common.utils import get_logger from common.exceptions import M2MReverseNotAllowed @@ -19,7 +19,11 @@ from .models import AssetPermission, RemoteAppPermission, RebuildUserTreeTask logger = get_logger(__file__) -# Todo: 检查授权规则到期，从而修改授权规则 +@receiver([pre_save], sender=AssetPermission) +def on_asset_perm_deactive(instance: AssetPermission, **kwargs): + old = AssetPermission.objects.only('is_active').get(id=instance.id) + if instance.is_active != old.is_active: + create_rebuild_user_tree_task_by_asset_perm(instance) @receiver([pre_delete], sender=AssetPermission) @@ -32,16 +36,17 @@ def create_rebuild_user_tree_task(user_ids): RebuildUserTreeTask.objects.bulk_create( [RebuildUserTreeTask(user_id=i) for i in user_ids] ) - transaction.on_commit(dispatch_mapping_node_tasks) + transaction.on_commit(dispatch_mapping_node_tasks.delay) def create_rebuild_user_tree_task_by_asset_perm(asset_perm: AssetPermission): - user_ap_query_name = AssetPermission.users.field.related_query_name() - group_ap_query_name = AssetPermission.user_groups.field.related_query_name() - - user_ap_q = Q(**{f'{user_ap_query_name}': asset_perm}) - group_ap_q = Q(**{f'groups__{group_ap_query_name}': asset_perm}) - user_ids = User.objects.filter(user_ap_q | group_ap_q).distinct().values_list('id', flat=True) + user_ids = set() + user_ids.update( + UserGroup.objects.filter(assetpermissions=asset_perm).distinct().values_list('users__id', flat=True) + ) + user_ids.update( + User.objects.filter(assetpermissions=asset_perm).distinct().values_list('id', flat=True) + ) create_rebuild_user_tree_task(user_ids) diff --git a/apps/perms/tasks.py b/apps/perms/tasks.py index 973553dc4..f427b684c 100644 --- a/apps/perms/tasks.py +++ b/apps/perms/tasks.py @@ -1,18 +1,22 @@ # ~*~ coding: utf-8 ~*~ from __future__ import absolute_import, unicode_literals +from datetime import timedelta +from django.db.models import Q +from django.conf import settings from celery import shared_task from common.utils import get_logger +from common.utils.timezone import now from users.models import User -from perms.models import RebuildUserTreeTask -from perms.utils.user_node_tree import rebuild_user_mapping_nodes_if_need_with_lock +from perms.models import RebuildUserTreeTask, AssetPermission +from perms.utils.user_asset_permission import rebuild_user_mapping_nodes_if_need_with_lock logger = get_logger(__file__) @shared_task(queue='node_tree') def rebuild_user_mapping_nodes_celery_task(user_id): - logger.info(f'rebuild user[{user_id}] mapping nodes') + logger.info(f'>>> rebuild user[{user_id}] mapping nodes') user = User.objects.get(id=user_id) rebuild_user_mapping_nodes_if_need_with_lock(user) @@ -21,5 +25,33 @@ def rebuild_user_mapping_nodes_celery_task(user_id): def dispatch_mapping_node_tasks(): user_ids = RebuildUserTreeTask.objects.all().values_list('user_id', flat=True).distinct() for id in user_ids: - logger.info(f'dispatch mapping node task for user[{id}]') + logger.info(f'>>> dispatch mapping node task for user[{id}]') rebuild_user_mapping_nodes_celery_task.delay(id) + + +@shared_task(queue='check_asset_perm_expired') +def check_asset_permission_expired(): + """ + 这里的任务要足够短，不要影响周期任务 + """ + periodic = settings.PERM_EXPIRED_CHECK_PERIODIC + end = now() + start = end - timedelta(seconds=periodic * 1.2) + ids = AssetPermission.objects.filter( + date_expired__gt=start, date_expired__lt=end + ).distinct().values_list('id', flat=True) + logger.info(f'>>> checking {start} to {end} have {ids} expired') + dispatch_process_expired_asset_permission.delay(ids) + + +@shared_task(queue='node_tree') +def dispatch_process_expired_asset_permission(asset_perm_ids): + user_ids = User.objects.filter( + Q(assetpermissions__id__in=asset_perm_ids) | + Q(groups__assetpermissions__id__in=asset_perm_ids) + ).distinct().values_list('id', flat=True) + RebuildUserTreeTask.objects.bulk_create( + [RebuildUserTreeTask(user_id=user_id) for user_id in user_ids] + ) + + dispatch_mapping_node_tasks.delay() diff --git a/apps/perms/urls/asset_permission.py b/apps/perms/urls/asset_permission.py index 43970af87..b266a0fb8 100644 --- a/apps/perms/urls/asset_permission.py +++ b/apps/perms/urls/asset_permission.py @@ -57,14 +57,22 @@ user_permission_urlpatterns = [ # 普通用户 -> 命令执行 -> 左侧树 path('nodes-with-assets/tree/', api.MyGrantedNodesWithAssetsAsTreeApi.as_view(), name='my-nodes-with-assets-as-tree'), - # Node children with assets as tree + # 主要用于 luna 页面，带资产的节点树 path('/nodes/children-with-assets/tree/', api.UserGrantedNodeChildrenWithAssetsAsTreeForAdminApi.as_view(), name='user-nodes-children-with-assets-as-tree'), path('nodes/children-with-assets/tree/', api.MyGrantedNodeChildrenWithAssetsAsTreeApi.as_view(), name='my-nodes-children-with-assets-as-tree'), - # Node assets + # 查询授权树上某个节点的所有资产 path('/nodes//assets/', api.UserGrantedNodeAssetsForAdminApi.as_view(), name='user-node-assets'), path('nodes//assets/', api.MyGrantedNodeAssetsApi.as_view(), name='my-node-assets'), + # 未分组的资产 + path('/nodes/ungrouped/assets/', api.UserDirectGrantedAssetsForAdminApi.as_view(), name='user-ungrouped-assets'), + path('nodes/ungrouped/assets/', api.MyDirectGrantedAssetsApi.as_view(), name='my-ungrouped-assets'), + + # 收藏的资产 + path('/nodes/favorite/assets/', api.UserFavoriteGrantedAssetsForAdminApi.as_view(), name='user-ungrouped-assets'), + path('nodes/favorite/assets/', api.MyFavoriteGrantedAssetsApi.as_view(), name='my-ungrouped-assets'), + # Asset System users path('/assets//system-users/', api.UserGrantedAssetSystemUsersForAdminApi.as_view(), name='user-asset-system-users'), path('assets//system-users/', api.MyGrantedAssetSystemUsersApi.as_view(), name='my-asset-system-users'), diff --git a/apps/perms/utils/__init__.py b/apps/perms/utils/__init__.py index 5ff0a8dfc..5ef34b6d9 100644 --- a/apps/perms/utils/__init__.py +++ b/apps/perms/utils/__init__.py @@ -5,4 +5,4 @@ from .asset_permission import * from .remote_app_permission import * from .database_app_permission import * from .k8s_app_permission import * -from .user_node_tree import * +from .user_asset_permission import * diff --git a/apps/perms/utils/asset_permission.py b/apps/perms/utils/asset_permission.py index 66c997895..6474b4334 100644 --- a/apps/perms/utils/asset_permission.py +++ b/apps/perms/utils/asset_permission.py @@ -4,47 +4,12 @@ from django.db.models import Q from common.utils import get_logger from ..models import AssetPermission -from ..hands import Asset, User -from users.models import UserGroup +from ..hands import Asset, User, UserGroup from perms.models.base import BasePermissionQuerySet logger = get_logger(__file__) -def get_user_permissions(user, include_group=True): - if include_group: - groups = user.groups.all() - arg = Q(users=user) | Q(user_groups__in=groups) - else: - arg = Q(users=user) - return AssetPermission.get_queryset_with_prefetch().filter(arg) - - -def get_user_group_permissions(user_group): - return AssetPermission.get_queryset_with_prefetch().filter( - user_groups=user_group - ) - - -def get_asset_permissions(asset, include_node=True): - if include_node: - nodes = asset.get_all_nodes(flat=True) - arg = Q(assets=asset) | Q(nodes__in=nodes) - else: - arg = Q(assets=asset) - return AssetPermission.objects.valid().filter(arg) - - -def get_node_permissions(node): - return AssetPermission.objects.valid().filter(nodes=node) - - -def get_system_user_permissions(system_user): - return AssetPermission.objects.valid().filter( - system_users=system_user - ) - - def get_asset_system_users_id_with_actions(asset_perm_queryset: BasePermissionQuerySet, asset: Asset): nodes = asset.get_nodes() node_keys = set() diff --git a/apps/perms/utils/user_node_tree.py b/apps/perms/utils/user_asset_permission.py similarity index 58% rename from apps/perms/utils/user_node_tree.py rename to apps/perms/utils/user_asset_permission.py index 8612bc3e5..80e4778ea 100644 --- a/apps/perms/utils/user_node_tree.py +++ b/apps/perms/utils/user_asset_permission.py @@ -6,15 +6,16 @@ import inspect from django.conf import settings from django.db.models import F, Q, Value, BooleanField +from django.utils.translation import gettext as _ from common.utils import get_logger from common.const.distributed_lock_key import UPDATE_MAPPING_NODE_TASK_LOCK_KEY from orgs.utils import tmp_to_root_org from common.utils.timezone import dt_formater, now -from assets.models import Node, Asset +from assets.models import Node, Asset, FavoriteAsset from django.db.transaction import atomic from orgs import lock -from perms.models import UserGrantedMappingNode, RebuildUserTreeTask +from perms.models import UserGrantedMappingNode, RebuildUserTreeTask, AssetPermission from users.models import User logger = get_logger(__name__) @@ -22,24 +23,35 @@ logger = get_logger(__name__) ADD = 'add' REMOVE = 'remove' +UNGROUPED_NODE_KEY = 'ungrouped' +FAVORITE_NODE_KEY = 'favorite' + +TMP_GRANTED_FIELD = '_granted' +TMP_ASSET_GRANTED_FIELD = '_asset_granted' +TMP_GRANTED_ASSETS_AMOUNT_FIELD = '_granted_assets_amount' + # 使用场景 -# Asset.objects.filter(get_granted_q(user)) -def get_granted_q(user: User): +# Asset.objects.filter(get_user_resources_q_granted_by_permissions(user)) +def get_user_resources_q_granted_by_permissions(user: User): + """ + 获取用户关联的 asset permission 或者 用户组关联的 asset permission 获取规则, + 前提 AssetPermission 对象中的 related_name 为 granted_by_permissions + :param user: + :return: + """ _now = now() return reduce(and_, ( Q(granted_by_permissions__date_start__lt=_now), Q(granted_by_permissions__date_expired__gt=_now), Q(granted_by_permissions__is_active=True), - (Q(granted_by_permissions__users=user) | Q(granted_by_permissions__user_groups__users=user)) + ( + Q(granted_by_permissions__users=user) | + Q(granted_by_permissions__user_groups__users=user) + ) )) -TMP_GRANTED_FIELD = '_granted' -TMP_ASSET_GRANTED_FIELD = '_asset_granted' -TMP_GRANTED_ASSETS_AMOUNT_FIELD = '_granted_assets_amount' - - # 使用场景 # `Node.objects.annotate(**node_annotate_mapping_node)` node_annotate_mapping_node = { @@ -56,7 +68,7 @@ node_annotate_set_granted = { } -def is_granted(node): +def is_direct_granted_by_annotate(node): return getattr(node, TMP_GRANTED_FIELD, False) @@ -139,7 +151,7 @@ def compute_tmp_mapping_node_from_perm(user: User): # 查询直接授权节点 nodes = Node.objects.filter( - get_granted_q(user) + get_user_resources_q_granted_by_permissions(user) ).distinct().only(*node_only_fields) granted_key_set = {_node.key for _node in nodes} @@ -165,7 +177,7 @@ def compute_tmp_mapping_node_from_perm(user: User): def process_direct_granted_assets(): # 查询直接授权资产 asset_ids = Asset.objects.filter( - get_granted_q(user) + get_user_resources_q_granted_by_permissions(user) ).distinct().values_list('id', flat=True) # 查询授权资产关联的节点设置 granted_asset_nodes = Node.objects.filter( @@ -227,7 +239,10 @@ def set_node_granted_assets_amount(user, node): if _granted: assets_amount = node.assets_amount else: - assets_amount = count_node_all_granted_assets(user, node.key) + if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: + assets_amount = count_direct_granted_node_assets(user, node.key) + else: + assets_amount = count_node_all_granted_assets(user, node.key) setattr(node, TMP_GRANTED_ASSETS_AMOUNT_FIELD, assets_amount) @@ -238,6 +253,68 @@ def rebuild_user_mapping_nodes(user): create_mapping_nodes(user, tmp_nodes) +def get_user_granted_nodes_list_via_mapping_node(user): + """ + 这里的 granted nodes, 是整棵树需要的node，推算出来的也算 + :param user: + :return: + """ + # 获取 `UserGrantedMappingNode` 中对应的 `Node` + nodes = Node.objects.filter( + mapping_nodes__user=user, + ).annotate( + **node_annotate_mapping_node + ).distinct() + + key_to_node_mapper = {} + nodes_descendant_q = Q() + + for node in nodes: + if not is_direct_granted_by_annotate(node): + # 未授权的节点资产数量设置为 `UserGrantedMappingNode` 中的数量 + node.assets_amount = get_granted_assets_amount(node) + else: + # 直接授权的节点 + # 增加查询后代节点的过滤条件 + nodes_descendant_q |= Q(key__startswith=f'{node.key}:') + key_to_node_mapper[node.key] = node + + if nodes_descendant_q: + descendant_nodes = Node.objects.filter( + nodes_descendant_q + ).annotate( + **node_annotate_set_granted + ) + for node in descendant_nodes: + key_to_node_mapper[node.key] = node + + all_nodes = key_to_node_mapper.values() + return all_nodes + + +def get_user_granted_all_assets(user, via_mapping_node=True): + asset_perm_ids = get_user_all_assetpermission_ids(user) + if via_mapping_node: + granted_node_keys = UserGrantedMappingNode.objects.filter( + user=user, granted=True, + ).values_list('key', flat=True).distinct() + else: + granted_node_keys = Node.objects.filter( + granted_by_permissions__id__in=asset_perm_ids + ).distinct().values_list('key', flat=True) + granted_node_keys = Node.clean_children_keys(granted_node_keys) + + granted_node_q = Q() + for _key in granted_node_keys: + granted_node_q |= Q(nodes__key__startswith=f'{_key}:') + granted_node_q |= Q(nodes__key=_key) + + assets__id = get_user_direct_granted_assets(user, asset_perm_ids).values_list('id', flat=True) + + q = granted_node_q | Q(id__in=list(assets__id)) + return Asset.org_objects.filter(q).distinct() + + def get_node_all_granted_assets(user: User, key): """ 此算法依据 `UserGrantedMappingNode` 的数据查询 @@ -249,9 +326,10 @@ def get_node_all_granted_assets(user: User, key): # 查询该节点下的授权节点 granted_mapping_nodes = UserGrantedMappingNode.objects.filter( - user=user, - granted=True, - ).filter(Q(key__startswith=f'{key}:') | Q(key=key)) + user=user, granted=True, + ).filter( + Q(key__startswith=f'{key}:') | Q(key=key) + ) # 根据授权节点构建资产查询条件 granted_nodes_qs = [] @@ -277,7 +355,7 @@ def get_node_all_granted_assets(user: User, key): if only_asset_granted_nodes_qs: only_asset_granted_nodes_q = reduce(or_, only_asset_granted_nodes_qs) - only_asset_granted_nodes_q &= get_granted_q(user) + only_asset_granted_nodes_q &= get_user_resources_q_granted_by_permissions(user) q.append(only_asset_granted_nodes_q) if q: @@ -285,36 +363,57 @@ def get_node_all_granted_assets(user: User, key): return assets +def get_direct_granted_node_ids(user: User, key): + granted_q = get_user_resources_q_granted_by_permissions(user) + + # 先查出该节点下的直接授权节点 + granted_nodes = Node.objects.filter( + Q(key__startswith=f'{key}:') | Q(key=key) + ).filter(granted_q).distinct().only('id', 'key') + + node_ids = set() + # 根据直接授权节点查询他们的子节点 + q = Q() + for _node in granted_nodes: + q |= Q(key__startswith=f'{_node.key}:') + node_ids.add(_node.id) + + if q: + descendant_ids = Node.objects.filter(q).values_list('id', flat=True).distinct() + node_ids.update(descendant_ids) + return node_ids + + def get_node_all_granted_assets_from_perm(user: User, key): """ 此算法依据 `AssetPermission` 的数据查询 1. 查询该节点下的直接授权节点 2. 查询该节点下授权资产关联的节点 """ - granted_q = get_granted_q(user) - - granted_nodes = Node.objects.filter( - Q(key__startswith=f'{key}:') | Q(key=key) - ).filter(granted_q).distinct() - + granted_q = get_user_resources_q_granted_by_permissions(user) # 直接授权资产查询条件 - granted_asset_filter_q = (Q(nodes__key__startswith=f'{key}:') | Q(nodes__key=key)) & granted_q - - # 根据授权节点构建资产查询条件 - q = granted_asset_filter_q - for _node in granted_nodes: - q |= Q(nodes__key__startswith=f'{_node.key}:') - q |= Q(nodes__key=_node.key) - + q = (Q(nodes__key__startswith=f'{key}:') | Q(nodes__key=key)) & granted_q + node_ids = get_direct_granted_node_ids(user, key) + q |= Q(nodes__id__in=node_ids) asset_qs = Asset.objects.filter(q).distinct() return asset_qs +def get_direct_granted_node_assets_from_perm(user: User, key): + node_ids = get_direct_granted_node_ids(user, key) + asset_qs = Asset.objects.filter(nodes__id__in=node_ids).distinct() + return asset_qs + + def count_node_all_granted_assets(user: User, key): return get_node_all_granted_assets_from_perm(user, key).count() -def get_ungranted_node_children(user, key=''): +def count_direct_granted_node_assets(user: User, key): + return get_direct_granted_node_assets_from_perm(user, key).count() + + +def get_indirect_granted_node_children(user, key=''): """ 获取用户授权树中未授权节点的子节点 只匹配在 `UserGrantedMappingNode` 中存在的节点 @@ -329,6 +428,72 @@ def get_ungranted_node_children(user, key=''): # 设置节点授权资产数量 for _node in nodes: - if not is_granted(_node): + if not is_direct_granted_by_annotate(_node): _node.assets_amount = get_granted_assets_amount(_node) return nodes + + +def get_top_level_granted_nodes(user): + nodes = list(get_indirect_granted_node_children(user, key='')) + if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: + ungrouped_node = get_ungrouped_node(user) + nodes.insert(0, ungrouped_node) + favorite_node = get_favorite_node(user) + nodes.insert(0, favorite_node) + return nodes + + +def get_user_all_assetpermission_ids(user: User): + asset_perm_ids = set() + asset_perm_ids.update( + AssetPermission.objects.valid().filter(users=user).distinct().values_list('id', flat=True) + ) + asset_perm_ids.update( + AssetPermission.objects.valid().filter(user_groups__users=user).distinct().values_list('id', flat=True) + ) + return asset_perm_ids + + +def get_user_direct_granted_assets(user, asset_perm_ids=None): + if asset_perm_ids is None: + asset_perm_ids = get_user_all_assetpermission_ids(user) + assets = Asset.org_objects.filter(granted_by_permissions__id__in=asset_perm_ids).distinct() + return assets + + +def count_user_direct_granted_assets(user): + count = get_user_direct_granted_assets(user).values_list('id').count() + return count + + +def get_ungrouped_node(user): + assets_amount = count_user_direct_granted_assets(user) + return Node( + id=UNGROUPED_NODE_KEY, + key=UNGROUPED_NODE_KEY, + value=_(UNGROUPED_NODE_KEY), + assets_amount=assets_amount + ) + + +def get_favorite_node(user): + assets_amount = FavoriteAsset.get_user_favorite_assets(user).values_list('id').count() + return Node( + id=FAVORITE_NODE_KEY, + key=FAVORITE_NODE_KEY, + value=_(FAVORITE_NODE_KEY), + assets_amount=assets_amount + ) + + +def init_user_tree_if_need(user): + """ + 升级授权树策略后，用户的数据可能还未初始化，为防止用户显示没有数据 + 先检查 MappingNode 如果没有数据，同步创建用户授权树 + """ + if not UserGrantedMappingNode.objects.filter(user=user).exists(): + try: + rebuild_user_mapping_nodes_with_lock(user) + except lock.SomeoneIsDoingThis: + # 您的数据正在初始化，请稍等 + raise lock.SomeoneIsDoingThis(detail=_('Please wait while your data is being initialized')) diff --git a/jms b/jms index 48eca674a..9c9fa92c3 100755 --- a/jms +++ b/jms @@ -158,6 +158,7 @@ def parse_service(s): all_services = [ 'gunicorn', 'celery_ansible', 'celery_default', 'beat', 'flower', 'daphne', 'celery_node_tree', + 'check_asset_perm_expired', ] if s == 'all': return all_services @@ -168,7 +169,7 @@ def parse_service(s): elif s == "task": return ["celery_ansible", "celery_default", "beat"] elif s == "celery": - return ["celery_ansible", "celery_default", "celery_node_tree"] + return ["celery_ansible", "celery_default", "celery_node_tree", "check_asset_perm_expired"] elif "," in s: services = set() for i in s.split(','): @@ -225,6 +226,11 @@ def get_start_celery_node_tree_kwargs(): return get_start_worker_kwargs('node_tree', 10) +def get_start_celery_check_asset_perm_expired_kwargs(): + print("\n- Start Celery as Distributed Task Queue: CheckAseetPermissionExpired") + return get_start_worker_kwargs('check_asset_perm_expired', 1) + + def get_start_worker_kwargs(queue, num): # Todo: Must set this environment, otherwise not no ansible result return os.environ.setdefault('PYTHONOPTIMIZE', '1') @@ -369,6 +375,7 @@ def start_service(s): "celery_ansible": get_start_celery_ansible_kwargs, "celery_default": get_start_celery_default_kwargs, "celery_node_tree": get_start_celery_node_tree_kwargs, + "check_asset_perm_expired": get_start_celery_check_asset_perm_expired_kwargs, "beat": get_start_beat_kwargs, "flower": get_start_flower_kwargs, "daphne": get_start_daphne_kwargs,