Commit a9d15381 authored by ibuler's avatar ibuler

[Update] 修改一些方法

parent f8ff223f
......@@ -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:
......
......@@ -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
......@@ -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):
"""
......
......@@ -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()
......
......@@ -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")
# 重新关联系统用户和资产的关系
......
......@@ -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:
......
# 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 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 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 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 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 get_user_group_permissions(user_group):
return AssetPermission.objects.all().valid().filter(
user_groups=user_group
)
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)
class AssetPermissionUtil:
@staticmethod
def get_user_permissions(user):
return AssetPermission.objects.all().valid().filter(users=user)
@staticmethod
def get_user_group_permissions(user_group):
return AssetPermission.objects.all().valid().filter(
user_groups=user_group
)
def get_node_permissions(node):
return AssetPermission.objects.all().valid().filter(nodes=node)
@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)
def get_system_user_permissions(system_user):
return AssetPermission.objects.valid().all().filter(
system_users=system_user
)
@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
class AssetPermissionUtil:
get_permissions_map = {
"User": get_user_permissions,
"UserGroup": get_user_group_permissions,
"Asset": get_asset_permissions,
"Node": get_node_permissions,
"SystemUser": get_node_permissions,
}
@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
def __init__(self, obj):
self.object = obj
self._permissions = None
@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
@property
def permissions(self):
if self._permissions:
return self._permissions
object_cls = self.object.__class__.__name__
func = self.get_permissions_map[object_cls]
permissions = func(self.object)
self._permissions = permissions
return permissions
@classmethod
def get_user_group_nodes_with_assets(cls, group):
def get_nodes_direct(self):
"""
:param group:
:return: {node: {asset: set(su1, su2)}}
返回用户/组授权规则直接关联的节点
:return: {node1: set(system_user1,)}
"""
_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)
permissions = self.permissions.prefetch_related('nodes', 'system_users')
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))
for node in perm.nodes.all():
nodes[node].update(perm.system_users.all())
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):
def get_assets_direct(self):
"""
:param user:
:return: {node: {asset: set(su1, su2)}}
返回用户授权规则直接关联的资产
:return: {asset1: set(system_user1,)}
"""
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)
assets = defaultdict(set)
permissions = self.permissions.prefetch_related('assets', 'system_users')
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()))
for asset in perm.assets.all().valid().prefetch_related('nodes'):
assets[asset].update(perm.system_users.all())
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
)
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
@property
def permissions(self):
if self._permissions:
return self._permissions
if isinstance(self.object, User):
pass
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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment