jumpserver/apps/perms/utils/asset_permission.py

# coding: utf-8

import time
import uuid
from collections import defaultdict
import json
from hashlib import md5
import itertools

from django.utils import timezone
from django.db.models import Q
from django.core.cache import cache
from django.conf import settings

from orgs.utils import set_to_root_org
from common.utils import get_logger, timeit
from common.tree import TreeNode
from assets.utils import NodeUtil
from .. import const
from ..models import AssetPermission, Action
from ..hands import Node, Asset
from .stack import PermSystemUserNodeUtil, PermAssetsAmountUtil

logger = get_logger(__file__)


__all__ = [
    'AssetPermissionUtil', 'is_obj_attr_has', 'sort_assets',
    'ParserNode',
]


class GenerateTree:
    def __init__(self):
        """
        nodes = {
          node.key: {
            "system_users": {
              system_user.id: actions,
            },
            "assets": set([asset.id,]),
          },
        }
        assets = {
           asset.id: {
             system_user.id: actions,
           },
        }
        """
        self._node_util = None
        self.nodes = defaultdict(lambda: {
            "system_users": defaultdict(int), "assets": set(),
            "assets_amount": 0, "all_assets": set(),
        })
        self.assets = defaultdict(lambda: defaultdict(int))
        self._root_node = None
        self._ungroup_node = None
        self._nodes_with_assets = None
        self._all_assets_nodes_key = None
        self._asset_counter = 0
        self._system_user_counter = 0
        self._nodes_assets_counter = 0

    @property
    def node_util(self):
        if not self._node_util:
            self._node_util = NodeUtil()
        return self._node_util

    @staticmethod
    def key_sort(key):
        key_list = [int(i) for i in key.split(':')]
        return len(key_list), key_list

    @property
    def root_key(self):
        if self._root_node:
            return self._root_node
        all_keys = self.nodes.keys()
        # 如果没有授权节点，就放到默认的根节点下
        if not all_keys:
            return None
        root_key = min(all_keys, key=self.key_sort)
        self._root_key = root_key
        return root_key

    @property
    def all_assets_nodes_keys(self):
        if not self._all_assets_nodes_key:
            self._all_assets_nodes_key = Asset.get_all_nodes_keys()
        return self._all_assets_nodes_key

    @property
    def ungrouped_key(self):
        if self._ungroup_node:
            return self._ungroup_node
        if self.root_key:
            node_key = "{}:{}".format(self.root_key, '-1')
        else:
            node_key = '1:-1'
        self._ungroup_node = node_key
        return node_key

    @timeit
    def add_assets_without_system_users(self, assets_ids):
        for asset_id in assets_ids:
            self.add_asset(asset_id, {})

    @timeit
    def add_assets(self, assets_ids_with_system_users):
        for asset_id, system_users_ids in assets_ids_with_system_users.items():
            self.add_asset(asset_id, system_users_ids)

    # @timeit
    def add_asset(self, asset_id, system_users_ids=None):
        """
        :param asset_id:
        :param system_users_ids: {system_user.id: actions, }
        :return:
        """
        if not system_users_ids:
            system_users_ids = defaultdict(int)

        # 获取已有资产的系统用户和actions，并更新到最新系统用户信息中
        old_system_users_ids = self.assets[asset_id]
        for system_user_id, action in old_system_users_ids.items():
            system_users_ids[system_user_id] |= action

        asset_nodes_keys = self.all_assets_nodes_keys.get(asset_id, [])
        # {asset.id: [node.key, ], }
        # 获取用户在的节点
        in_nodes = set(self.nodes.keys()) & set(asset_nodes_keys)
        if not in_nodes:
            self.nodes[self.ungrouped_key]["assets"].add(asset_id)
            self.assets[asset_id] = system_users_ids
            return

        # 遍历用户应该在的节点
        for key in in_nodes:
            # 把自己加入到树上的节点中
            self.nodes[key]["assets"].add(asset_id)
            # 获取自己所在节点的系统用户，并添加进去
            node_system_users_ids = self.nodes[key]["system_users"]
            for system_user_id, action in node_system_users_ids.items():
                system_users_ids[system_user_id] |= action
        self.assets[asset_id] = system_users_ids

    def add_node(self, node_key, system_users_ids=None):
        """
        :param node_key: node.key
        :param system_users_ids: {system_user.id: actions,}
        :return:
        """
        if not system_users_ids:
            system_users_ids = defaultdict(int)
        self.nodes[node_key]["system_users"] = system_users_ids

    # 添加树节点
    @timeit
    def add_nodes(self, nodes_keys_with_system_users_ids):
        """
        :param nodes_keys_with_system_users_ids:
        {node.key: {system_user.id: actions,}, }
        :return:
        """
        util = PermSystemUserNodeUtil()
        family = util.get_nodes_family_and_system_users(nodes_keys_with_system_users_ids)
        for key, system_users in family.items():
            self.add_node(key, system_users)

    def get_assets(self):
        """
        :return:
        [
            {"id": asset.id, "system_users": {system_user.id: actions, }},
        ]
        """
        assets = []
        for asset_id, system_users in self.assets.items():
            assets.append({"id": asset_id, "system_users": system_users})
        return assets

    def set_ungrouped_assets_nodes_if_need(self):
        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
            return
        ungrouped_assets_ids = self.nodes[self.ungrouped_key]["assets"]
        for asset_id in ungrouped_assets_ids:
            in_nodes = self.all_assets_nodes_keys.get(asset_id, [])
            for node_key in in_nodes:
                parents_keys = self.node_util.get_nodes_parents_keys_by_key(node_key, with_self=False)
                for parent_key in parents_keys:
                    n = self.nodes[parent_key]
                self.nodes[node_key]["assets"].add(asset_id)
        self.nodes.pop(self.ungrouped_key, None)

    @timeit
    def get_nodes_with_assets(self):
        """
        :return:
        [
            {
                'key': node.key,
                'assets_amount': 10
                'assets': {
                    asset.id: {
                        system_user.id: actions,
                    },
                },
            },
        ]
        """
        if self._nodes_with_assets:
            return self._nodes_with_assets
        self.set_ungrouped_assets_nodes_if_need()
        util = PermAssetsAmountUtil()
        nodes_with_assets_amount = util.compute_nodes_assets_amount(self.nodes)
        nodes = []
        for key, values in nodes_with_assets_amount.items():
            assets = {asset_id: self.assets.get(asset_id) for asset_id in values["assets"]}
            all_assets = {asset_id: self.assets.get(asset_id) for asset_id in values["all_assets"]}
            nodes.append({
                "key": key, "assets": assets, "all_assets": all_assets,
                "assets_amount": values["assets_amount"]
            })
        # 如果返回空节点，页面构造授权资产树报错
        if not nodes:
            nodes.append({
                "key": const.EMPTY_NODE_KEY, "assets": {}, "assets_amount": 0,
                "all_assets": {},
            })
        nodes.sort(key=lambda n: self.key_sort(n["key"]))
        self._nodes_with_assets = nodes
        return nodes

    def get_nodes(self):
        self.set_ungrouped_assets_nodes_if_need()
        nodes = list(self.nodes.keys())
        if not nodes:
            nodes.append(const.EMPTY_NODE_KEY)
        return list(nodes)


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 AssetPermissionCacheMixin:
    CACHE_KEY_PREFIX = '_ASSET_PERM_CACHE_V2_'
    CACHE_META_KEY_PREFIX = '_ASSET_PERM_META_KEY_V2_'
    CACHE_TIME = settings.ASSETS_PERM_CACHE_TIME
    CACHE_POLICY_MAP = (('0', 'never'), ('1', 'using'), ('2', 'refresh'))
    cache_policy = '1'
    obj_id = ''
    _filter_id = None

    @classmethod
    def is_not_using_cache(cls, cache_policy):
        return cls.CACHE_TIME == 0 or cache_policy in cls.CACHE_POLICY_MAP[0]

    @classmethod
    def is_using_cache(cls, cache_policy):
        return cls.CACHE_TIME != 0 and cache_policy in cls.CACHE_POLICY_MAP[1]

    @classmethod
    def is_refresh_cache(cls, cache_policy):
        return cache_policy in cls.CACHE_POLICY_MAP[2]

    def _is_not_using_cache(self):
        return self.is_not_using_cache(self.cache_policy)

    def _is_using_cache(self):
        return self.is_using_cache(self.cache_policy)

    def _is_refresh_cache(self):
        return self.is_refresh_cache(self.cache_policy)

    #@timeit
    def get_cache_key(self, resource):
        cache_key = self.CACHE_KEY_PREFIX + '{obj_id}_{filter_id}_{resource}'
        return cache_key.format(
            obj_id=self.obj_id, filter_id=self._filter_id,
            resource=resource
        )

    @property
    def node_asset_key(self):
        return self.get_cache_key('NODES_WITH_ASSETS')

    @property
    def node_key(self):
        return self.get_cache_key('NODES')

    @property
    def asset_key(self):
        key = self.get_cache_key('ASSETS')
        return key

    @property
    def system_key(self):
        return self.get_cache_key('SYSTEM_USER')

    def get_resource_from_cache(self, resource):
        logger.debug("Try get resource from cache")
        key_map = {
            "assets": self.asset_key,
            "nodes": self.node_key,
            "nodes_with_assets": self.node_asset_key,
            "system_users": self.system_key
        }
        key = key_map.get(resource)
        if not key:
            raise ValueError("Not a valid resource: {}".format(resource))
        cached = cache.get(key)
        if not cached:
            logger.debug("Not found resource cache, update it")
            self.update_cache()
            cached = cache.get(key)
        return cached

    def get_resource(self, resource):
        if self._is_using_cache():
            logger.debug("Using cache to get resource")
            return self.get_resource_from_cache(resource)
        elif self._is_refresh_cache():
            logger.debug("Need refresh cache")
            self.expire_cache()
            data = self.get_resource_from_cache(resource)
            return data
        else:
            logger.debug("Not using cache get source")
            return self.get_resource_without_cache(resource)

    def get_resource_without_cache(self, resource):
        attr = 'get_{}_without_cache'.format(resource)
        return getattr(self, attr)()

    def get_nodes_with_assets(self):
        return self.get_resource("nodes_with_assets")

    def get_assets(self):
        return self.get_resource("assets")

    def get_nodes(self):
        return self.get_resource("nodes")

    def get_system_users(self):
        return self.get_resource("system_users")

    def get_meta_cache_key(self):
        cache_key = self.CACHE_META_KEY_PREFIX + '{obj_id}_{filter_id}'
        key = cache_key.format(
            obj_id=self.obj_id, filter_id=self._filter_id
        )
        return key

    @property
    def cache_meta(self):
        key = self.get_meta_cache_key()
        meta = cache.get(key) or {}
        # print("Meta key: {}".format(key))
        # print("Meta id: {}".format(meta["id"]))
        return meta

    def update_cache(self):
        assets = self.get_resource_without_cache("assets")
        nodes_with_assets = self.get_resource_without_cache("nodes_with_assets")
        system_users = self.get_resource_without_cache("system_users")
        nodes = self.get_resource_without_cache("nodes")
        cache.set(self.asset_key, assets, self.CACHE_TIME)
        cache.set(self.node_asset_key, nodes_with_assets, self.CACHE_TIME)
        cache.set(self.system_key, system_users, self.CACHE_TIME)
        cache.set(self.node_key, nodes, self.CACHE_TIME)
        self.set_meta_to_cache()

    def set_meta_to_cache(self):
        key = self.get_meta_cache_key()
        meta = {
            'id': str(uuid.uuid4()),
            'datetime': timezone.now(),
            'object': str(self.object)
        }
        # print("Set meta key: {}".format(key))
        # print("set meta to cache: {}".format(meta["id"]))
        cache.set(key, meta, self.CACHE_TIME)

    def expire_cache_meta(self):
        cache_key = self.CACHE_META_KEY_PREFIX + '{obj_id}_*'
        key = cache_key.format(obj_id=self.obj_id)
        cache.delete_pattern(key)

    def expire_cache(self):
        """
        因为 获取用户的节点，资产，系统用户等都能会缓存，这里会清理所有与该对象有关的
        缓存，以免造成不统一的情况
        :return:
        """
        cache_key = self.CACHE_KEY_PREFIX + '{obj_id}_*'
        key = cache_key.format(obj_id='*')
        cache.delete_pattern(key)
        self.expire_cache_meta()

    @classmethod
    def expire_all_cache(cls):
        key = cls.CACHE_KEY_PREFIX + '*'
        cache.delete_pattern(key)
        meta_key = cls.CACHE_META_KEY_PREFIX + '*'
        cache.delete_pattern(meta_key)


class AssetPermissionUtil(AssetPermissionCacheMixin):
    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, cache_policy='0'):
        self.object = obj
        self.obj_id = str(obj.id)
        self._permissions = None
        self._permissions_id = None  # 标记_permission的唯一值
        self._assets = None
        self._filter_id = 'None'  # 当通过filter更改 permission是标记
        self.cache_policy = cache_policy
        self.tree = GenerateTree()
        self.change_org_if_need()
        self.nodes = None
        self._nodes = None
        self._assets_direct = None
        self._nodes_direct = None
        self.node_util = NodeUtil()
        self.tree._node_util = self.node_util

    @staticmethod
    def change_org_if_need():
        set_to_root_org()

    @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

    @timeit
    def filter_permissions(self, **filters):
        filters_json = json.dumps(filters, sort_keys=True)
        self._permissions = self.permissions.filter(**filters)
        self._filter_id = md5(filters_json.encode()).hexdigest()

    @timeit
    def get_nodes_direct(self):
        """
        返回直接授权的节点，
        并将节点添加到tree.nodes中，并将节点下的资产添加到tree.assets中
        :return:
        {node.key: {system_user.id: actions,}, }
        """
        if self._nodes_direct:
            return self._nodes_direct
        nodes_keys = defaultdict(lambda: defaultdict(int))
        for perm in self.permissions:
            actions = [perm.actions]
            system_users_ids = [s.id for s in perm.system_users.all()]
            _nodes_keys = [n.key for n in perm.nodes.all()]
            iterable = itertools.product(_nodes_keys, system_users_ids, actions)
            for node_key, sys_id, action in iterable:
                nodes_keys[node_key][sys_id] |= action

        self.tree.add_nodes(nodes_keys)

        all_nodes_keys = set()
        for key in nodes_keys:
            children_keys = self.node_util.get_all_children_keys_by_key(key)
            all_nodes_keys.update(set(children_keys))

        if all_nodes_keys:
            assets_ids = Asset.objects.filter(
                nodes__key__in=all_nodes_keys
            ).valid().values_list("id", flat=True).distinct()
        else:
            assets_ids = []
        self.tree.add_assets_without_system_users(assets_ids)
        self._nodes_direct = nodes_keys
        return nodes_keys

    def get_nodes_without_cache(self):
        self.get_assets_without_cache()
        return self.tree.get_nodes()

    @timeit
    def get_assets_direct(self):
        """
        返回直接授权的资产，
        并添加到tree.assets中
        :return:
        {asset.id: {system_user.id: actions, }, }
        """
        if self._assets_direct:
            return self._assets_direct
        assets_ids = defaultdict(lambda: defaultdict(int))
        for perm in self.permissions:
            actions = [perm.actions]
            _assets_ids = perm.assets.valid().values_list("id", flat=True)
            system_users_ids = perm.system_users.values_list("id", flat=True)
            iterable = itertools.product(_assets_ids, system_users_ids, actions)
            for asset_id, sys_id, action in iterable:
                assets_ids[asset_id][sys_id] |= action
        self.tree.add_assets(assets_ids)
        self._assets_direct = assets_ids
        return assets_ids

    @timeit
    def get_assets_without_cache(self):
        """
        :return:
        [
            {"id": asset.id, "system_users": {system_user.id: actions, }},
        ]
        """
        if self._assets:
            return self._assets
        self.get_nodes_direct()
        self.get_assets_direct()
        assets = self.tree.get_assets()
        self._assets = assets
        return assets

    @timeit
    def get_nodes_with_assets_without_cache(self):
        self.get_assets_without_cache()
        nodes_assets = self.tree.get_nodes_with_assets()
        return nodes_assets

    def get_system_users_without_cache(self):
        system_users = set()
        permissions = self.permissions.prefetch_related('system_users')
        for perm in permissions:
            system_users.update(perm.system_users.all())
        return system_users


def is_obj_attr_has(obj, val, attrs=("hostname", "ip", "comment")):
    if not attrs:
        vals = [val for val in obj.__dict__.values() if isinstance(val, (str, int))]
    else:
        vals = [getattr(obj, attr) for attr in attrs if
                hasattr(obj, attr) and isinstance(hasattr(obj, attr), (str, int))]

    for v in vals:
        if str(v).find(val) != -1:
            return True
    return False


def sort_assets(assets, order_by='hostname', reverse=False):
    if order_by == 'ip':
        assets = sorted(assets, key=lambda asset: [int(d) for d in asset.ip.split('.') if d.isdigit()], reverse=reverse)
    else:
        assets = sorted(assets, key=lambda asset: getattr(asset, order_by), reverse=reverse)
    return assets


class ParserNode:
    nodes_only_fields = ("key", "value", "id")
    assets_only_fields = ("platform", "hostname", "id", "ip", "protocols")
    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_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, system_users):
        icon_skin = 'file'
        if asset.platform.lower() == 'windows':
            icon_skin = 'windows'
        elif asset.platform.lower() == 'linux':
            icon_skin = 'linux'
        _system_users = []
        for system_user in system_users:
            _system_users.append({
                'id': system_user.id,
                'name': system_user.name,
                'username': system_user.username,
                'protocol': system_user.protocol,
                'priority': system_user.priority,
                'login_mode': system_user.login_mode,
                'actions': [Action.value_to_choices(system_user.actions)],
            })
        data = {
            'id': str(asset.id),
            'name': asset.hostname,
            'title': asset.ip,
            'pId': node.key,
            'isParent': False,
            'open': False,
            'iconSkin': icon_skin,
            'meta': {
                'system_users': _system_users,
                'type': 'asset',
                'asset': {
                    'id': asset.id,
                    'hostname': asset.hostname,
                    'ip': asset.ip,
                    'protocols': asset.protocols_as_list,
                    'platform': asset.platform,
                },
            }
        }
        tree_node = TreeNode(**data)
        return tree_node