jumpserver/apps/perms/tree.py

from django.utils.translation import gettext_lazy as _
from django.db.models import Q, TextChoices
from assets.tree.node_tree import AssetNodeTree
from assets.tree.type_tree import AssetTypeTree
from assets.models import Asset, FavoriteAsset
from perms.utils.user_perm_asset import UserPermAssetUtil
from common.utils import lazyproperty

__all__ = ['UserAssetNodeTree', 'UserAssetTypeTree']


class UserAssetNodeTree(AssetNodeTree):

    class SpecialKey(TextChoices):
        FAVORITE = 'favorite', _('Favorite')
        UNGROUPED = 'ungrouped', _('Ungrouped')

    def __init__(self, user, org=None):
        super().__init__(org=org)
        self.user = user
        self.org = org
        self.util = UserPermAssetUtil(user=user, org=org)
        self.cache_key_scope_assets_ids = 'cache_key_org_{}_user_{}_scope_assets_ids'.format(
            str(self.org.id), str(self.util._user.id)
        )
    
    def _scope_assets_ids(self):
        return self.util.user_permed_all_assets_ids

    def get_node(self, key):
        node = super().get_node(key)
        if node:
            return node
        if key == self.SpecialKey.FAVORITE:
            return self.favorite_tree_node
        if key == self.SpecialKey.UNGROUPED:
            return self.ungrouped_tree_node
    
    def get_extra_nodes(self):
        f_tree_node = self.create_favorite_tree_node()
        u_tree_node = self.create_ungrouped_tree_node()
        return [f_tree_node, u_tree_node]
    
    def create_favorite_tree_node(self):
        self.favorite_tree_node.assets_amount_total = self.favorite_assets_ids.count()
        return self.favorite_tree_node

    @lazyproperty
    def favorite_tree_node(self):
        kwargs = {
            'raw_id': self.SpecialKey.FAVORITE,
            'key': self.SpecialKey.FAVORITE,
            'name': self.SpecialKey.FAVORITE.label,
            'parent_key': None
        }
        tree_node = self.create_tree_node(**kwargs)
        return tree_node

    def create_ungrouped_tree_node(self):
        self.ungrouped_tree_node.assets_amount_total = len(self.ungrouped_assets_ids)
        return self.ungrouped_tree_node

    @lazyproperty
    def ungrouped_tree_node(self):
        kwargs = {
            'raw_id': self.SpecialKey.UNGROUPED,
            'key': self.SpecialKey.UNGROUPED,
            'name': self.SpecialKey.UNGROUPED.label,
            'parent_key': None
        }
        tree_node = self.create_tree_node(**kwargs)
        return tree_node
    
    def _filter_node_assets(self, node_key):
        if node_key == self.SpecialKey.FAVORITE:
            return self.filter_favorite_assets()
        if node_key == self.SpecialKey.UNGROUPED:
            return self.filter_ungrouped_assets()
        return super()._filter_node_assets(node_key)
    
    @property
    def ungrouped_assets_ids(self):
        return self.util._user_direct_asset_ids

    @property
    def favorite_assets_ids(self):
        return FavoriteAsset.get_user_favorite_asset_ids(user=self.user)

    def filter_favorite_assets(self):
        return Asset.objects.filter(id__in=self.favorite_assets_ids)

    def filter_ungrouped_assets(self):
        return Asset.objects.filter(id__in=self.util._user_direct_asset_ids)


class UserAssetTypeTree(AssetTypeTree):

    def __init__(self, user, org=None):
        super().__init__(org=org)
        self.user = user
        self.util = UserPermAssetUtil(user=user, org=org)
        self.cache_key_platform_assets_amount_mapper = f'''
            cache_key_user_{str(self.user.id)}_org_{str(self.org.id)}_platform_assets_amount_mapper 
        '''
    
    @lazyproperty
    def scope_assets_queryset(self):
        assets = self.util.user_permed_all_assets()
        return assets