mirror of https://github.com/jumpserver/jumpserver
				
				
				
			
		
			
				
	
	
		
			485 lines
		
	
	
		
			16 KiB
		
	
	
	
		
			Python
		
	
	
			
		
		
	
	
			485 lines
		
	
	
		
			16 KiB
		
	
	
	
		
			Python
		
	
	
# coding: utf-8
 | 
						||
import time
 | 
						||
import pickle
 | 
						||
from collections import defaultdict
 | 
						||
from functools import reduce
 | 
						||
 | 
						||
from django.core.cache import cache
 | 
						||
from django.db.models import Q
 | 
						||
from django.conf import settings
 | 
						||
 | 
						||
from orgs.utils import current_org
 | 
						||
from common.utils import get_logger, timeit, lazyproperty
 | 
						||
from common.tree import TreeNode
 | 
						||
from assets.utils import TreeService
 | 
						||
from ..models import AssetPermission
 | 
						||
from ..hands import Node, Asset, SystemUser, User, FavoriteAsset
 | 
						||
 | 
						||
logger = get_logger(__file__)
 | 
						||
 | 
						||
 | 
						||
__all__ = [
 | 
						||
    'ParserNode', 'AssetPermissionUtil',
 | 
						||
]
 | 
						||
 | 
						||
 | 
						||
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
 | 
						||
    )
 | 
						||
 | 
						||
 | 
						||
class AssetPermissionUtilCacheMixin:
 | 
						||
    user_tree_cache_key = 'USER_PERM_TREE_{}_{}_{}'
 | 
						||
    user_tree_cache_ttl = settings.ASSETS_PERM_CACHE_TIME
 | 
						||
    user_tree_cache_enable = settings.ASSETS_PERM_CACHE_ENABLE
 | 
						||
    user_tree_map = {}
 | 
						||
    cache_policy = '0'
 | 
						||
    obj_id = ''
 | 
						||
    _filter_id = 'None'
 | 
						||
 | 
						||
    @property
 | 
						||
    def cache_key(self):
 | 
						||
        return self.get_cache_key()
 | 
						||
 | 
						||
    def get_cache_key(self, org_id=None):
 | 
						||
        if org_id is None:
 | 
						||
            org_id = current_org.org_id()
 | 
						||
 | 
						||
        key = self.user_tree_cache_key.format(
 | 
						||
            org_id, self.obj_id, self._filter_id
 | 
						||
        )
 | 
						||
        return key
 | 
						||
 | 
						||
    def expire_user_tree_cache(self):
 | 
						||
        cache.delete(self.cache_key)
 | 
						||
 | 
						||
    @classmethod
 | 
						||
    def expire_all_user_tree_cache(cls):
 | 
						||
        expire_cache_key = "USER_TREE_EXPIRED_AT"
 | 
						||
        latest_expired = cache.get(expire_cache_key, 0)
 | 
						||
        now = time.time()
 | 
						||
        if now - latest_expired < 60:
 | 
						||
            return
 | 
						||
        key = cls.user_tree_cache_key.format('*', '1', '1')
 | 
						||
        key = key.replace('_1', '')
 | 
						||
        cache.delete_pattern(key)
 | 
						||
        cache.set(expire_cache_key, now)
 | 
						||
 | 
						||
    @classmethod
 | 
						||
    def expire_org_tree_cache(cls, org_id=None):
 | 
						||
        if org_id is None:
 | 
						||
            org_id = current_org.org_id()
 | 
						||
        key = cls.user_tree_cache_key.format(org_id, '*', '1')
 | 
						||
        key = key.replace('_1', '')
 | 
						||
        cache.delete_pattern(key)
 | 
						||
 | 
						||
    def set_user_tree_to_cache(self, user_tree):
 | 
						||
        data = pickle.dumps(user_tree)
 | 
						||
        cache.set(self.cache_key, data, self.user_tree_cache_ttl)
 | 
						||
 | 
						||
    def get_user_tree_from_cache(self):
 | 
						||
        data = cache.get(self.cache_key)
 | 
						||
        if not data:
 | 
						||
            return None
 | 
						||
        user_tree = pickle.loads(data)
 | 
						||
        return user_tree
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def get_user_tree_from_cache_if_need(self):
 | 
						||
        if not self.user_tree_cache_enable:
 | 
						||
            return None
 | 
						||
        if self.cache_policy == '1':
 | 
						||
            return self.get_user_tree_from_cache()
 | 
						||
        elif self.cache_policy == '2':
 | 
						||
            self.expire_user_tree_cache()
 | 
						||
            return None
 | 
						||
        else:
 | 
						||
            return None
 | 
						||
 | 
						||
    def set_user_tree_to_cache_if_need(self, user_tree):
 | 
						||
        if self.cache_policy == '0':
 | 
						||
            return
 | 
						||
        if not self.user_tree_cache_enable:
 | 
						||
            return None
 | 
						||
        self.set_user_tree_to_cache(user_tree)
 | 
						||
 | 
						||
 | 
						||
class AssetPermissionUtil(AssetPermissionUtilCacheMixin):
 | 
						||
    get_permissions_map = {
 | 
						||
        "User": get_user_permissions,
 | 
						||
        "UserGroup": get_user_group_permissions,
 | 
						||
        "Asset": get_asset_permissions,
 | 
						||
        "Node": get_node_permissions,
 | 
						||
        "SystemUser": get_system_user_permissions,
 | 
						||
    }
 | 
						||
    assets_only = (
 | 
						||
        'id', 'hostname', 'ip', "platform", "domain_id",
 | 
						||
        'comment', 'is_active', 'os', 'org_id'
 | 
						||
    )
 | 
						||
 | 
						||
    def __init__(self, obj=None, cache_policy='0'):
 | 
						||
        self.object = obj
 | 
						||
        self.cache_policy = cache_policy
 | 
						||
        self.obj_id = str(obj.id) if obj else None
 | 
						||
        self._permissions = None
 | 
						||
        self._filter_id = 'None'  # 当通过filter更改 permission是标记
 | 
						||
        self.change_org_if_need()
 | 
						||
        self._user_tree = None
 | 
						||
        self._user_tree_filter_id = 'None'
 | 
						||
 | 
						||
        if not isinstance(obj, User):
 | 
						||
            self.cache_policy = '0'
 | 
						||
 | 
						||
    @staticmethod
 | 
						||
    def change_org_if_need():
 | 
						||
        pass
 | 
						||
 | 
						||
    @lazyproperty
 | 
						||
    def full_tree(self):
 | 
						||
        return Node.tree()
 | 
						||
 | 
						||
    @property
 | 
						||
    def permissions(self):
 | 
						||
        if self._permissions:
 | 
						||
            return self._permissions
 | 
						||
        if self.object is None:
 | 
						||
            return AssetPermission.objects.none()
 | 
						||
        object_cls = self.object.__class__.__name__
 | 
						||
        func = self.get_permissions_map[object_cls]
 | 
						||
        permissions = func(self.object)
 | 
						||
        self._permissions = permissions
 | 
						||
        return permissions
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def filter_permissions(self, **filters):
 | 
						||
        self.cache_policy = '0'
 | 
						||
        self._permissions = self.permissions.filter(**filters)
 | 
						||
 | 
						||
    @lazyproperty
 | 
						||
    def user_tree(self):
 | 
						||
        return self.get_user_tree()
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def get_assets_direct(self):
 | 
						||
        """
 | 
						||
        返回直接授权的资产,
 | 
						||
        并添加到tree.assets中
 | 
						||
        :return:
 | 
						||
        {asset.id: {system_user.id: actions, }, }
 | 
						||
        """
 | 
						||
        assets_ids = self.permissions.values_list('assets', flat=True)
 | 
						||
        return Asset.objects.filter(id__in=assets_ids)
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def get_nodes_direct(self):
 | 
						||
        """
 | 
						||
        返回直接授权的节点,
 | 
						||
        并将节点添加到tree.nodes中,并将节点下的资产添加到tree.assets中
 | 
						||
        :return:
 | 
						||
        {node.key: {system_user.id: actions,}, }
 | 
						||
        """
 | 
						||
        nodes_ids = self.permissions.values_list('nodes', flat=True)
 | 
						||
        return Node.objects.filter(id__in=nodes_ids)
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def add_direct_nodes_to_user_tree(self, user_tree):
 | 
						||
        """
 | 
						||
        将授权规则的节点放到用户树上, 从full tree中粘贴子树
 | 
						||
        """
 | 
						||
        nodes_direct_keys = self.permissions \
 | 
						||
            .exclude(nodes__isnull=True) \
 | 
						||
            .values_list('nodes__key', flat=True) \
 | 
						||
            .distinct()
 | 
						||
        nodes_direct_keys = list(nodes_direct_keys)
 | 
						||
        # 排序,保证从上层节点开始加
 | 
						||
        nodes_direct_keys.sort(key=lambda x: len(x))
 | 
						||
        for key in nodes_direct_keys:
 | 
						||
            # 如果树上已经有这个节点,代表子树已经存在
 | 
						||
            if user_tree.contains(key):
 | 
						||
                continue
 | 
						||
            # 找到这个节点的父节点,如果父节点不在树上,则挂到ROOT上
 | 
						||
            parent = self.full_tree.parent(key)
 | 
						||
            if not user_tree.contains(parent.identifier):
 | 
						||
                parent = user_tree.root_node()
 | 
						||
            subtree = self.full_tree.subtree(key)
 | 
						||
            user_tree.paste(parent.identifier, subtree, deep=True)
 | 
						||
 | 
						||
        for node in user_tree.all_nodes_itr():
 | 
						||
            assets = list(self.full_tree.assets(node.identifier))
 | 
						||
            user_tree.set_assets(node.identifier, assets)
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def add_single_assets_node_to_user_tree(self, user_tree):
 | 
						||
        """
 | 
						||
        将单独授权的资产放到树上,如果设置了单独资产到 未分组中,则放到未分组中
 | 
						||
        如果没有,则查询资产属于的资产组,放到树上
 | 
						||
        """
 | 
						||
        # 添加单独授权资产的节点
 | 
						||
        nodes_single_assets = defaultdict(set)
 | 
						||
        queryset = self.permissions.exclude(assets__isnull=True) \
 | 
						||
            .values_list('assets', 'assets__nodes__key') \
 | 
						||
            .distinct()
 | 
						||
 | 
						||
        for item in queryset:
 | 
						||
            nodes_single_assets[item[1]].add(item[0])
 | 
						||
        nodes_single_assets.pop(None, None)
 | 
						||
 | 
						||
        for key in tuple(nodes_single_assets.keys()):
 | 
						||
            if user_tree.contains(key):
 | 
						||
                nodes_single_assets.pop(key)
 | 
						||
 | 
						||
        if not nodes_single_assets:
 | 
						||
            return
 | 
						||
 | 
						||
        # 如果要设置到ungroup中
 | 
						||
        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
 | 
						||
            node_key = Node.ungrouped_key
 | 
						||
            node_value = Node.ungrouped_value
 | 
						||
            user_tree.create_node(
 | 
						||
                identifier=node_key, tag=node_value,
 | 
						||
                parent=user_tree.root,
 | 
						||
            )
 | 
						||
            assets = set()
 | 
						||
            for _assets in nodes_single_assets.values():
 | 
						||
                assets.update(set(_assets))
 | 
						||
            user_tree.set_assets(node_key, assets)
 | 
						||
            return
 | 
						||
 | 
						||
        # 获取单独授权资产,并没有在授权的节点上
 | 
						||
        for key, assets in nodes_single_assets.items():
 | 
						||
            if not self.full_tree.contains(key):
 | 
						||
                continue
 | 
						||
            node = self.full_tree.get_node(key, deep=True)
 | 
						||
            parent_id = self.full_tree.parent(key).identifier
 | 
						||
            parent = user_tree.get_node(parent_id)
 | 
						||
            if not parent:
 | 
						||
                parent = user_tree.root_node()
 | 
						||
            user_tree.add_node(node, parent)
 | 
						||
            user_tree.set_assets(node.identifier, assets)
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def parse_user_tree_to_full_tree(self, user_tree):
 | 
						||
        """
 | 
						||
        经过前面两个动作,用户授权的节点已放到树上,但是树不是完整的,
 | 
						||
        这里要讲树构造成一个完整的书
 | 
						||
        """
 | 
						||
        # 开始修正user_tree,保证父节点都在树上
 | 
						||
        root_children = user_tree.children('')
 | 
						||
        for child in root_children:
 | 
						||
            if child.identifier.isdigit():
 | 
						||
                continue
 | 
						||
            if child.identifier.startswith('-'):
 | 
						||
                continue
 | 
						||
            ancestors = self.full_tree.ancestors(
 | 
						||
                child.identifier, with_self=False, deep=True
 | 
						||
            )
 | 
						||
            if not ancestors:
 | 
						||
                continue
 | 
						||
            user_tree.safe_add_ancestors(child, ancestors)
 | 
						||
 | 
						||
    def add_favorite_node_if_need(self, user_tree):
 | 
						||
        if not isinstance(self.object, User):
 | 
						||
            return
 | 
						||
        node_key = Node.favorite_key
 | 
						||
        node_value = Node.favorite_value
 | 
						||
        user_tree.create_node(
 | 
						||
            identifier=node_key, tag=node_value,
 | 
						||
            parent=user_tree.root,
 | 
						||
        )
 | 
						||
        node = user_tree.get_node(node_key)
 | 
						||
        assets_id = FavoriteAsset.get_user_favorite_assets_id(self.object)
 | 
						||
        all_valid_assets = user_tree.all_valid_assets(user_tree.root)
 | 
						||
        valid_assets_id = set(assets_id) & all_valid_assets
 | 
						||
        user_tree.set_assets(node_key, valid_assets_id)
 | 
						||
        # 必须设置这个,否则看不到个数
 | 
						||
        node.data['all_assets'] = None
 | 
						||
 | 
						||
    def set_user_tree_to_local(self, user_tree):
 | 
						||
        self._user_tree = user_tree
 | 
						||
        self._user_tree_filter_id = self._filter_id
 | 
						||
 | 
						||
    def get_user_tree_from_local(self):
 | 
						||
        if self._user_tree and self._user_tree_filter_id == self._filter_id:
 | 
						||
            return self._user_tree
 | 
						||
        return None
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def get_user_tree(self):
 | 
						||
        user_tree = self.get_user_tree_from_cache_if_need()
 | 
						||
        if user_tree:
 | 
						||
            return user_tree
 | 
						||
        user_tree = TreeService()
 | 
						||
        full_tree_root = self.full_tree.root_node()
 | 
						||
        user_tree.create_node(
 | 
						||
            tag=full_tree_root.tag,
 | 
						||
            identifier=full_tree_root.identifier
 | 
						||
        )
 | 
						||
        self.add_direct_nodes_to_user_tree(user_tree)
 | 
						||
        self.add_single_assets_node_to_user_tree(user_tree)
 | 
						||
        self.parse_user_tree_to_full_tree(user_tree)
 | 
						||
        self.add_favorite_node_if_need(user_tree)
 | 
						||
        self.set_user_tree_to_cache_if_need(user_tree)
 | 
						||
        self.set_user_tree_to_local(user_tree)
 | 
						||
        return user_tree
 | 
						||
 | 
						||
    # Todo: 是否可以获取多个资产的系统用户
 | 
						||
    def get_asset_system_users_id_with_actions(self, asset):
 | 
						||
        nodes = asset.get_nodes()
 | 
						||
        nodes_keys_related = set()
 | 
						||
        for node in nodes:
 | 
						||
            ancestor_keys = node.get_ancestor_keys(with_self=True)
 | 
						||
            nodes_keys_related.update(set(ancestor_keys))
 | 
						||
        kwargs = {"assets": asset}
 | 
						||
 | 
						||
        if nodes_keys_related:
 | 
						||
            kwargs["nodes__key__in"] = nodes_keys_related
 | 
						||
 | 
						||
        queryset = self.permissions
 | 
						||
        if kwargs == 1:
 | 
						||
            queryset = queryset.filter(**kwargs)
 | 
						||
        elif len(kwargs) > 1:
 | 
						||
            kwargs = [{k: v} for k, v in kwargs.items()]
 | 
						||
            args = [Q(**kw) for kw in kwargs]
 | 
						||
            args = reduce(lambda x, y: x | y, args)
 | 
						||
            queryset = queryset.filter(args)
 | 
						||
        else:
 | 
						||
            queryset = queryset.none()
 | 
						||
        asset_protocols = asset.protocols_as_dict.keys()
 | 
						||
        values = queryset.filter(system_users__protocol__in=asset_protocols).distinct()\
 | 
						||
            .values_list('system_users', 'actions')
 | 
						||
        system_users_actions = defaultdict(int)
 | 
						||
        for system_user_id, actions in values:
 | 
						||
            if None in (system_user_id, actions):
 | 
						||
                continue
 | 
						||
            for i, action in values:
 | 
						||
                system_users_actions[i] |= actions
 | 
						||
        return system_users_actions
 | 
						||
 | 
						||
    def get_permissions_nodes_and_assets(self):
 | 
						||
        from assets.models import Node
 | 
						||
        permissions = self.permissions
 | 
						||
        nodes_keys = permissions.exclude(nodes__isnull=True)\
 | 
						||
            .values_list('nodes__key', flat=True)
 | 
						||
        assets_ids = permissions.exclude(assets__isnull=True)\
 | 
						||
            .values_list('assets', flat=True)
 | 
						||
        nodes_keys = set(nodes_keys)
 | 
						||
        assets_ids = set(assets_ids)
 | 
						||
        nodes_keys = Node.clean_children_keys(nodes_keys)
 | 
						||
        return nodes_keys, assets_ids
 | 
						||
 | 
						||
    @timeit
 | 
						||
    def get_assets(self):
 | 
						||
        nodes_keys, assets_ids = self.get_permissions_nodes_and_assets()
 | 
						||
        queryset = Node.get_nodes_all_assets(
 | 
						||
            nodes_keys, extra_assets_ids=assets_ids
 | 
						||
        )
 | 
						||
        return queryset.valid()
 | 
						||
 | 
						||
    def get_nodes_assets(self, node, deep=False):
 | 
						||
        if deep:
 | 
						||
            assets_ids = self.user_tree.all_assets(node.key)
 | 
						||
        else:
 | 
						||
            assets_ids = self.user_tree.assets(node.key)
 | 
						||
        queryset = Asset.objects.filter(id__in=assets_ids)
 | 
						||
        return queryset.valid()
 | 
						||
 | 
						||
    def get_nodes(self):
 | 
						||
        return [n.identifier for n in self.user_tree.all_nodes_itr()]
 | 
						||
 | 
						||
    def get_system_users(self):
 | 
						||
        system_users_id = self.permissions.values_list('system_users', flat=True).distinct()
 | 
						||
        return SystemUser.objects.filter(id__in=system_users_id)
 | 
						||
 | 
						||
 | 
						||
class ParserNode:
 | 
						||
    nodes_only_fields = ("key", "value", "id")
 | 
						||
    assets_only_fields = ("hostname", "id", "ip", "protocols", "domain", "org_id")
 | 
						||
    system_users_only_fields = (
 | 
						||
        "id", "name", "username", "protocol", "priority", "login_mode",
 | 
						||
    )
 | 
						||
 | 
						||
    @staticmethod
 | 
						||
    def parse_node_to_tree_node(node):
 | 
						||
        name = '{} ({})'.format(node.value, node.assets_amount)
 | 
						||
        data = {
 | 
						||
            'id': node.key,
 | 
						||
            'name': name,
 | 
						||
            'title': name,
 | 
						||
            'pId': node.parent_key,
 | 
						||
            'isParent': True,
 | 
						||
            'open': node.is_org_root(),
 | 
						||
            'meta': {
 | 
						||
                'node': {
 | 
						||
                    "id": node.id,
 | 
						||
                    "key": node.key,
 | 
						||
                    "value": node.value,
 | 
						||
                },
 | 
						||
                'type': 'node'
 | 
						||
            }
 | 
						||
        }
 | 
						||
        tree_node = TreeNode(**data)
 | 
						||
        return tree_node
 | 
						||
 | 
						||
    @staticmethod
 | 
						||
    def parse_asset_to_tree_node(node, asset):
 | 
						||
        icon_skin = 'file'
 | 
						||
        platform = asset.platform_base.lower()
 | 
						||
        if platform == 'windows':
 | 
						||
            icon_skin = 'windows'
 | 
						||
        elif platform == 'linux':
 | 
						||
            icon_skin = 'linux'
 | 
						||
        parent_id = node.key if node else ''
 | 
						||
        data = {
 | 
						||
            'id': str(asset.id),
 | 
						||
            'name': asset.hostname,
 | 
						||
            'title': asset.ip,
 | 
						||
            'pId': parent_id,
 | 
						||
            'isParent': False,
 | 
						||
            'open': False,
 | 
						||
            'iconSkin': icon_skin,
 | 
						||
            'nocheck': not asset.has_protocol('ssh'),
 | 
						||
            'meta': {
 | 
						||
                'type': 'asset',
 | 
						||
                'asset': {
 | 
						||
                    'id': asset.id,
 | 
						||
                    'hostname': asset.hostname,
 | 
						||
                    'ip': asset.ip,
 | 
						||
                    'protocols': asset.protocols_as_list,
 | 
						||
                    'platform': asset.platform_base,
 | 
						||
                    'domain': asset.domain_id,
 | 
						||
                    'org_name': asset.org_name,
 | 
						||
                },
 | 
						||
            }
 | 
						||
        }
 | 
						||
        tree_node = TreeNode(**data)
 | 
						||
        return tree_node
 |