diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index a6a1d48d0..a974d3385 100644 --- a/apps/assets/models/asset.py +++ b/apps/assets/models/asset.py @@ -152,29 +152,13 @@ class Asset(models.Model): def get_all_nodes(self, flat=False): nodes = [] - for node in self.get_nodes_or_cache(): + for node in self.get_nodes(): _nodes = node.get_ancestor(with_self=True) _nodes.append(_nodes) if flat: nodes = list(reduce(lambda x, y: set(x) | set(y), nodes)) return nodes - @property - def nodes_cache_key(self): - key = "NODES_OF_{}".format(str(self.id)) - return key - - def get_nodes_or_cache(self): - cached = cache.get(self.nodes_cache_key) - if cached is not None: - return cached - nodes = list(self.get_nodes()) - cache.set(self.nodes_cache_key, nodes, 3600) - return nodes - - def expire_nodes_cache(self): - cache.delete(self.nodes_cache_key) - @property def hardware_info(self): if self.cpu_count: diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index 88b685d28..dea6c01ef 100644 --- a/apps/assets/models/node.py +++ b/apps/assets/models/node.py @@ -22,6 +22,21 @@ class Node(models.Model): def __str__(self): return self.full_value + def __eq__(self, other): + return self.key == other.key + + def __gt__(self, other): + if self.is_root(): + return True + self_key = [int(k) for k in self.key.split(':')] + other_key = [int(k) for k in other.key.split(':')] + if len(self_key) < len(other_key): + return True + elif len(self_key) > len(other_key): + return False + else: + return self_key[-1] < other_key[-1] + @property def name(self): return self.value @@ -93,7 +108,7 @@ class Node(models.Model): Q(nodes__id=self.id) | Q(nodes__isnull=True) ) else: - assets = Asset.objects.filter(nodes__id=self.id) + assets = self.assets.all() return assets def get_valid_assets(self): @@ -104,8 +119,8 @@ class Node(models.Model): if self.is_root(): assets = Asset.objects.all() else: - nodes = self.get_all_children(with_self=True) - assets = Asset.objects.filter(nodes__in=nodes).distinct() + pattern = r'^{0}$|^{0}:'.format(self.key) + assets = Asset.objects.filter(nodes__key__regex=pattern) return assets def get_all_valid_assets(self): @@ -154,10 +169,3 @@ class Node(models.Model): return obj -class Tree: - def __init__(self, root): - self.root = root - self.nodes = [] - - def add_node(self, node): - pass diff --git a/apps/assets/serializers/asset.py b/apps/assets/serializers/asset.py index ac666e3a7..a0fdfab73 100644 --- a/apps/assets/serializers/asset.py +++ b/apps/assets/serializers/asset.py @@ -16,8 +16,6 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer): """ 资产的数据结构 """ - nodes = serializers.SerializerMethodField() - class Meta: model = Asset list_serializer_class = BulkListSerializer @@ -31,10 +29,6 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer): ]) return fields - @staticmethod - def get_nodes(obj): - return [n.id for n in obj.get_nodes_or_cache()] - class AssetGrantedSerializer(serializers.ModelSerializer): """ diff --git a/apps/assets/serializers/node.py b/apps/assets/serializers/node.py index 73639a100..90e140710 100644 --- a/apps/assets/serializers/node.py +++ b/apps/assets/serializers/node.py @@ -68,7 +68,10 @@ class NodeSerializer(serializers.ModelSerializer): @staticmethod def get_assets_amount(obj): - return obj.get_all_assets().count() + if obj.is_node: + return obj.get_all_assets().count() + else: + return 0 def get_fields(self): fields = super().get_fields() diff --git a/apps/assets/signals_handler.py b/apps/assets/signals_handler.py index 16459c786..157c88012 100644 --- a/apps/assets/signals_handler.py +++ b/apps/assets/signals_handler.py @@ -64,7 +64,6 @@ def on_system_user_assets_change(sender, instance=None, **kwargs): @receiver(m2m_changed, sender=Asset.nodes.through) def on_asset_node_changed(sender, instance=None, **kwargs): if isinstance(instance, Asset): - instance.expire_nodes_cache() if kwargs['action'] == 'post_add': logger.debug("Asset node change signal received") nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set']) @@ -81,10 +80,6 @@ def on_asset_node_changed(sender, instance=None, **kwargs): def on_node_assets_changed(sender, instance=None, **kwargs): if isinstance(instance, Node): assets = kwargs['model'].objects.filter(pk__in=kwargs['pk_set']) - # 清理资产节点缓存 - for asset in assets: - asset.expire_nodes_cache() - if kwargs['action'] == 'post_add': logger.debug("Node assets change signal received") # 重新关联系统用户和资产的关系 diff --git a/apps/perms/api.py b/apps/perms/api.py index b4172d03f..33a027064 100644 --- a/apps/perms/api.py +++ b/apps/perms/api.py @@ -70,7 +70,8 @@ class UserGrantedAssetsApi(ListAPIView): else: user = self.request.user - for k, v in AssetPermissionUtil.get_user_assets(user).items(): + util = AssetPermissionUtil(user) + for k, v in util.get_assets().items(): if k.is_unixlike(): system_users_granted = [s for s in v if s.protocol == 'ssh'] else: @@ -95,7 +96,8 @@ class UserGrantedNodesApi(ListAPIView): user = get_object_or_404(User, id=user_id) else: user = self.request.user - nodes = AssetPermissionUtil.get_user_nodes_with_assets(user) + util = AssetPermissionUtil(user) + nodes = util.get_nodes_with_assets() return nodes.keys() def get_permissions(self): @@ -116,7 +118,8 @@ class UserGrantedNodesWithAssetsApi(ListAPIView): else: user = get_object_or_404(User, id=user_id) - nodes = AssetPermissionUtil.get_user_nodes_with_assets(user) + util = AssetPermissionUtil(user) + nodes = util.get_nodes_with_assets() for node, _assets in nodes.items(): assets = _assets.keys() for k, v in _assets.items(): @@ -147,8 +150,9 @@ class UserGrantedNodeAssetsApi(ListAPIView): user = get_object_or_404(User, id=user_id) else: user = self.request.user + util = AssetPermissionUtil(user) node = get_object_or_404(Node, id=node_id) - nodes = AssetPermissionUtil.get_user_nodes_with_assets(user) + nodes = util.get_nodes_with_assets() assets = nodes.get(node, []) for asset, system_users in assets.items(): asset.system_users_granted = system_users @@ -172,7 +176,8 @@ class UserGroupGrantedAssetsApi(ListAPIView): return queryset user_group = get_object_or_404(UserGroup, id=user_group_id) - assets = AssetPermissionUtil.get_user_group_assets(user_group) + util = AssetPermissionUtil(user_group) + assets = util.get_assets() for k, v in assets.items(): k.system_users_granted = v queryset.append(k) @@ -189,7 +194,8 @@ class UserGroupGrantedNodesApi(ListAPIView): if group_id: group = get_object_or_404(UserGroup, id=group_id) - nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(group) + util = AssetPermissionUtil(group) + nodes = util.get_nodes_with_assets() return nodes.keys() return queryset @@ -206,7 +212,8 @@ class UserGroupGrantedNodesWithAssetsApi(ListAPIView): return queryset user_group = get_object_or_404(UserGroup, id=user_group_id) - nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(user_group) + util = AssetPermissionUtil(user_group) + nodes = util.get_nodes_with_assets() for node, _assets in nodes.items(): assets = _assets.keys() for asset, system_users in _assets.items(): @@ -226,7 +233,8 @@ class UserGroupGrantedNodeAssetsApi(ListAPIView): user_group = get_object_or_404(UserGroup, id=user_group_id) node = get_object_or_404(Node, id=node_id) - nodes = AssetPermissionUtil.get_user_group_nodes_with_assets(user_group) + util = AssetPermissionUtil(user_group) + nodes = util.get_nodes_with_assets() assets = nodes.get(node, []) for asset, system_users in assets.items(): asset.system_users_granted = system_users @@ -246,7 +254,8 @@ class ValidateUserAssetPermissionView(APIView): asset = get_object_or_404(Asset, id=asset_id) system_user = get_object_or_404(SystemUser, id=system_id) - assets_granted = AssetPermissionUtil.get_user_assets(user) + util = AssetPermissionUtil(user) + assets_granted = util.get_assets() if system_user in assets_granted.get(asset, []): return Response({'msg': True}, status=200) else: diff --git a/apps/perms/utils.py b/apps/perms/utils.py index 91f530f1d..36c432826 100644 --- a/apps/perms/utils.py +++ b/apps/perms/utils.py @@ -1,301 +1,118 @@ # coding: utf-8 from __future__ import absolute_import, unicode_literals -import collections from collections import defaultdict from django.db.models import Q -from django.utils import timezone -import copy -from common.utils import set_or_append_attr_bulk, get_logger +from common.utils import get_logger from .models import AssetPermission -from .hands import Node, User, UserGroup, Asset, SystemUser +from .hands import Node logger = get_logger(__file__) -class Tree: - def __init__(self): - self.__all_nodes = list(Node.objects.all().prefetch_related('assets')) - self.__node_asset_map = defaultdict(set) - self.nodes = defaultdict(dict) - self.root = Node.root() - self.init_node_asset_map() +def get_user_permissions(user, include_group=True): + if include_group: + groups = user.groups.all() + arg = Q(users=user) | Q(user_groups=groups) + else: + arg = Q(users=user) + return AssetPermission.objects.all().valid().filter(arg) - def init_node_asset_map(self): - for node in self.__all_nodes: - assets = node.get_assets().values_list('id', flat=True) - for asset in assets: - self.__node_asset_map[str(asset)].add(node) - def add_asset(self, asset, system_users): - nodes = self.__node_asset_map.get(str(asset.id), []) - self.add_nodes(nodes) - for node in nodes: - self.nodes[node][asset].update(system_users) +def get_user_group_permissions(user_group): + return AssetPermission.objects.all().valid().filter( + user_groups=user_group + ) - def add_node(self, node): - if node in self.nodes: - return - else: - self.nodes[node] = defaultdict(set) - if node.key == self.root.key: - return - parent_key = ':'.join(node.key.split(':')[:-1]) - for n in self.__all_nodes: - if n.key == parent_key: - self.add_node(n) - break - def add_nodes(self, nodes): - for node in nodes: - self.add_node(node) +def get_asset_permissions(asset, include_node=True): + if include_node: + nodes = asset.get_all_nodes(flat=True) + arg = Q(assets=asset) | Q(nodes=nodes) + else: + arg = Q(assets=asset) + return AssetPermission.objects.all().valid().filter(arg) + + +def get_node_permissions(node): + return AssetPermission.objects.all().valid().filter(nodes=node) + + +def get_system_user_permissions(system_user): + return AssetPermission.objects.valid().all().filter( + system_users=system_user + ) class AssetPermissionUtil: - @staticmethod - def get_user_permissions(user): - return AssetPermission.objects.all().valid().filter(users=user) + get_permissions_map = { + "User": get_user_permissions, + "UserGroup": get_user_group_permissions, + "Asset": get_asset_permissions, + "Node": get_node_permissions, + "SystemUser": get_node_permissions, + } - @staticmethod - def get_user_group_permissions(user_group): - return AssetPermission.objects.all().valid().filter( - user_groups=user_group - ) - - @staticmethod - def get_asset_permissions(asset): - return AssetPermission.objects.all().valid().filter( - assets=asset - ) - - @staticmethod - def get_node_permissions(node): - return AssetPermission.objects.all().valid().filter(nodes=node) - - @staticmethod - def get_system_user_permissions(system_user): - return AssetPermission.objects.valid().all().filter( - system_users=system_user - ) - - @classmethod - def get_user_group_nodes(cls, group): - nodes = defaultdict(set) - permissions = cls.get_user_group_permissions(group) - for perm in permissions: - _nodes = perm.nodes.all() - _system_users = perm.system_users.all() - set_or_append_attr_bulk(_nodes, 'permission', perm.id) - for node in _nodes: - nodes[node].update(set(_system_users)) - return nodes - - @classmethod - def get_user_group_assets_direct(cls, group): - assets = defaultdict(set) - permissions = cls.get_user_group_permissions(group) - for perm in permissions: - _assets = perm.assets.all().valid() - _system_users = perm.system_users.all() - set_or_append_attr_bulk(_assets, 'permission', perm.id) - for asset in _assets: - assets[asset].update(set(_system_users)) - return assets - - @classmethod - def get_user_group_nodes_assets(cls, group): - assets = defaultdict(set) - nodes = cls.get_user_group_nodes(group) - for node, _system_users in nodes.items(): - _assets = node.get_all_valid_assets() - set_or_append_attr_bulk(_assets, 'inherit_node', node.id) - set_or_append_attr_bulk(_assets, 'permission', getattr(node, 'permission', None)) - for asset in _assets: - assets[asset].update(set(_system_users)) - return assets - - @classmethod - def get_user_group_assets(cls, group): - assets = defaultdict(set) - _assets = cls.get_user_group_assets_direct(group) - _nodes_assets = cls.get_user_group_nodes_assets(group) - for asset, _system_users in _assets.items(): - assets[asset].update(set(_system_users)) - for asset, _system_users in _nodes_assets.items(): - assets[asset].update(set(_system_users)) - return assets - - @classmethod - def get_user_group_nodes_with_assets(cls, group): - """ - :param group: - :return: {node: {asset: set(su1, su2)}} - """ - _assets = cls.get_user_group_assets(group) - tree = Tree() - for asset, _system_users in _assets.items(): - _nodes = asset.get_nodes_or_cache() - tree.add_nodes(_nodes) - for node in _nodes: - tree.nodes[node][asset].update(_system_users) - return tree.nodes - - @classmethod - def get_user_assets_direct(cls, user): - assets = defaultdict(set) - permissions = list(cls.get_user_permissions(user)) - for perm in permissions: - _assets = perm.assets.all().valid() - _system_users = perm.system_users.all() - set_or_append_attr_bulk(_assets, 'permission', perm.id) - for asset in _assets: - assets[asset].update(set(_system_users)) - return assets - - @classmethod - def get_user_nodes_direct(cls, user): - nodes = defaultdict(set) - permissions = cls.get_user_permissions(user) - for perm in permissions: - _nodes = perm.nodes.all() - _system_users = perm.system_users.all() - set_or_append_attr_bulk(_nodes, 'permission', perm.id) - for node in _nodes: - nodes[node].update(set(_system_users)) - return nodes - - @classmethod - def get_user_nodes_inherit_group(cls, user): - nodes = defaultdict(set) - groups = user.groups.all() - for group in groups: - _nodes = cls.get_user_group_nodes(group) - for node, system_users in _nodes.items(): - nodes[node].update(set(system_users)) - return nodes - - @classmethod - def get_user_nodes(cls, user): - nodes = cls.get_user_nodes_direct(user) - nodes_inherit = cls.get_user_nodes_inherit_group(user) - for node, system_users in nodes_inherit.items(): - nodes[node].update(set(system_users)) - return nodes - - @classmethod - def get_user_nodes_assets_direct(cls, user): - assets = defaultdict(set) - nodes = cls.get_user_nodes_direct(user) - for node, _system_users in nodes.items(): - _assets = node.get_all_valid_assets() - set_or_append_attr_bulk(_assets, 'inherit_node', node.id) - set_or_append_attr_bulk(_assets, 'permission', getattr(node, 'permission', None)) - for asset in _assets: - assets[asset].update(set(_system_users)) - return assets - - @classmethod - def get_user_assets_inherit_group(cls, user): - assets = defaultdict(set) - for group in user.groups.all(): - _assets = cls.get_user_group_assets(group) - set_or_append_attr_bulk(_assets, 'inherit_group', group.id) - for asset, _system_users in _assets.items(): - assets[asset].update(_system_users) - return assets - - @classmethod - def get_user_assets(cls, user): - assets = defaultdict(set) - _assets_direct = cls.get_user_assets_direct(user) - _nodes_assets_direct = cls.get_user_nodes_assets_direct(user) - _assets_inherit_group = cls.get_user_assets_inherit_group(user) - for asset, _system_users in _assets_direct.items(): - assets[asset].update(_system_users) - for asset, _system_users in _nodes_assets_direct.items(): - assets[asset].update(_system_users) - for asset, _system_users in _assets_inherit_group.items(): - assets[asset].update(_system_users) - return assets - - @classmethod - def get_user_nodes_with_assets(cls, user): - """ - :param user: - :return: {node: {asset: set(su1, su2)}} - """ - tree = Tree() - _assets = cls.get_user_assets(user) - for asset, _system_users in _assets.items(): - tree.add_asset(asset, _system_users) - # _nodes = asset.get_nodes() - # tree.add_nodes(_nodes) - # for node in _nodes: - # tree.nodes[node][asset].update(_system_users) - return tree.nodes - - @classmethod - def get_system_user_assets(cls, system_user): - assets = set() - permissions = cls.get_system_user_permissions(system_user) - for perm in permissions: - assets.update(set(perm.assets.all().valid())) - nodes = perm.nodes.all() - for node in nodes: - assets.update(set(node.get_all_valid_assets())) - return assets - - @classmethod - def get_node_system_users(cls, node): - system_users = set() - permissions = cls.get_node_permissions(node) - for perm in permissions: - system_users.update(perm.system_users.all()) - return system_users - - -class AssetPermissionUtilsV2: def __init__(self, obj): self.object = obj self._permissions = None - @staticmethod - def get_user_permissions(user): - groups = user.groups.all() - return AssetPermission.objects.all().valid().filter( - Q(users=user) | Q(user_groups=groups) - ) - - @staticmethod - def get_user_group_permissions(user_group): - return AssetPermission.objects.all().valid().filter( - user_groups=user_group - ) - - @staticmethod - def get_asset_permissions(asset): - direct_nodes = asset.get_nodes_or_cache() - - return AssetPermission.objects.all().valid().filter( - Q(assets=asset) | Q(nodes=direct_nodes) - ) - - @staticmethod - def get_node_permissions(node): - return AssetPermission.objects.all().valid().filter(nodes=node) - - @staticmethod - def get_system_user_permissions(system_user): - return AssetPermission.objects.valid().all().filter( - system_users=system_user - ) - @property def permissions(self): if self._permissions: return self._permissions - if isinstance(self.object, User): - pass + object_cls = self.object.__class__.__name__ + func = self.get_permissions_map[object_cls] + permissions = func(self.object) + self._permissions = permissions + return permissions + + def get_nodes_direct(self): + """ + 返回用户/组授权规则直接关联的节点 + :return: {node1: set(system_user1,)} + """ + nodes = defaultdict(set) + permissions = self.permissions.prefetch_related('nodes', 'system_users') + for perm in permissions: + for node in perm.nodes.all(): + nodes[node].update(perm.system_users.all()) + return nodes + + def get_assets_direct(self): + """ + 返回用户授权规则直接关联的资产 + :return: {asset1: set(system_user1,)} + """ + assets = defaultdict(set) + permissions = self.permissions.prefetch_related('assets', 'system_users') + for perm in permissions: + for asset in perm.assets.all().valid().prefetch_related('nodes'): + assets[asset].update(perm.system_users.all()) + return assets + + def get_assets(self): + assets = self.get_assets_direct() + nodes = self.get_nodes_direct() + for node, system_users in nodes.items(): + _assets = node.get_all_assets().valid().prefetch_related('nodes') + for asset in _assets: + if isinstance(asset, Node): + print(_assets) + assets[asset].update(system_users) + return assets + + def get_nodes_with_assets(self): + assets = self.get_assets() + nodes = defaultdict(dict) + for asset, system_users in assets.items(): + _nodes = asset.nodes.all() + for node in _nodes: + if asset in nodes[node]: + nodes[node][asset].update(system_users) + else: + nodes[node] = {asset: system_users} + return nodes