jumpserver/apps/perms/utils.py

# coding: utf-8

from __future__ import absolute_import, unicode_literals
import uuid
from collections import defaultdict
from django.utils import timezone
from django.db.models import Q
from django.core.cache import cache
from django.conf import settings

from common.utils import get_logger
from common.tree import TreeNode
from .models import AssetPermission
from .hands import Node

logger = get_logger(__file__)


class GenerateTree:
    def __init__(self):
        """
        nodes: {"node_instance": {
            "asset_instance": set("system_user")
        }
        """
        self.__all_nodes = list(Node.objects.all())
        self.nodes = defaultdict(dict)

    def add_asset(self, asset, system_users):
        nodes = asset.nodes.all()
        self.add_nodes(nodes)
        for node in nodes:
            self.nodes[node][asset].update(system_users)

    def get_nodes(self):
        for node in self.nodes:
            assets = set(self.nodes.get(node).keys())
            for n in self.nodes.keys():
                if n.key.startswith(node.key + ':'):
                    assets.update(set(self.nodes[n].keys()))
            node.assets_amount = len(assets)
        return self.nodes

    def add_node(self, node):
        if node in self.nodes:
            return
        else:
            self.nodes[node] = defaultdict(set)
        if node.is_root():
            return
        for n in self.__all_nodes:
            if n.key == node.parent_key:
                self.add_node(n)
                break

    def add_nodes(self, nodes):
        for node in nodes:
            self.add_node(node)


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.objects.all().valid().filter(arg)


def get_user_group_permissions(user_group):
    return AssetPermission.objects.all().valid().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.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:
    get_permissions_map = {
        "User": get_user_permissions,
        "UserGroup": get_user_group_permissions,
        "Asset": get_asset_permissions,
        "Node": get_node_permissions,
        "SystemUser": get_node_permissions,
    }

    CACHE_KEY = '_ASSET_PERM_CACHE_{}_{}'
    CACHE_META_KEY = '_ASSET_PERM_META_KEY_{}'
    CACHE_TIME = settings.ASSETS_PERM_CACHE_TIME
    CACHE_POLICY_MAP = (('0', 'never'), ('1', 'using'), ('2', 'refresh'))

    def __init__(self, obj, cache_policy='0'):
        self.object = obj
        self.obj_id = str(obj.id)
        self._permissions = None
        self._assets = None
        self.cache_policy = cache_policy
        self.node_key = self.CACHE_KEY.format(self.obj_id, 'NODES_WITH_ASSETS')
        self.asset_key = self.CACHE_KEY.format(self.obj_id, 'ASSETS')
        self.system_key = self.CACHE_KEY.format(self.obj_id, 'SYSTEM_USER')

    @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

    def filter_permission_with_system_user(self, system_user):
        self._permissions = self.permissions.filter(system_users=system_user)

    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.filter(protocol=asset.protocol)
                )
        return assets

    def get_assets_without_cache(self):
        if self._assets:
            return self._assets
        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:
                assets[asset].update(
                    [s for s in system_users if s.protocol == asset.protocol]
                )
        self._assets = assets
        return self._assets

    def get_assets_from_cache(self):
        cached = cache.get(self.asset_key)
        if not cached:
            self.update_cache()
            cached = cache.get(self.asset_key)
        return cached

    def get_assets(self):
        if self.CACHE_TIME <= 0 or self.cache_policy in self.CACHE_POLICY_MAP[1]:
            return self.get_assets_from_cache()
        elif self.cache_policy in self.CACHE_POLICY_MAP[2]:
            self.expire_cache()
            return self.get_assets_from_cache()
        else:
            self.expire_cache()
            return self.get_assets_without_cache()

    def get_nodes_with_assets_without_cache(self):
        """
        返回节点并且包含资产
        {"node": {"assets": set("system_user")}}
        :return:
        """
        assets = self.get_assets_without_cache()
        tree = GenerateTree()
        for asset, system_users in assets.items():
            tree.add_asset(asset, system_users)
        return tree.get_nodes()

    def get_nodes_with_assets_from_cache(self):
        cached = cache.get(self.node_key)
        if not cached:
            self.update_cache()
            cached = cache.get(self.node_key)
        return cached

    def get_nodes_with_assets(self):
        if self.CACHE_TIME <= 0 or self.cache_policy in self.CACHE_POLICY_MAP[1]:
            return self.get_nodes_with_assets_from_cache()
        elif self.cache_policy in self.CACHE_POLICY_MAP[2]:
            self.expire_cache()
            return self.get_nodes_with_assets_from_cache()
        else:
            return self.get_nodes_with_assets_without_cache()

    def get_system_user_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 get_system_user_from_cache(self):
        cached = cache.get(self.system_key)
        if not cached:
            self.update_cache()
            cached = cache.get(self.system_key)
        return cached

    def get_system_users(self):
        if self.CACHE_TIME <= 0 or self.cache_policy in self.CACHE_POLICY_MAP[1]:
            return self.get_system_user_from_cache()
        elif self.cache_policy in self.CACHE_POLICY_MAP[2]:
            self.expire_cache()
            return self.get_system_user_from_cache()
        else:
            return self.get_system_user_without_cache()

    @property
    def cache_meta(self):
        key = self.CACHE_META_KEY.format(str(self.object.id))
        return cache.get(key) or {}

    def set_cache_meta(self):
        key = self.CACHE_META_KEY.format(str(self.object.id))
        meta = {
            'id': str(uuid.uuid4()),
            'datetime': timezone.now(),
            'object': str(self.object)
        }
        cache.set(key, meta, self.CACHE_TIME)

    def expire_cache_meta(self):
        key = self.CACHE_META_KEY.format(str(self.object.id))
        cache.delete(key)

    def update_cache(self):
        assets = self.get_assets_without_cache()
        nodes = self.get_nodes_with_assets_without_cache()
        system_users = self.get_system_user_without_cache()
        cache.set(self.asset_key, assets, self.CACHE_TIME)
        cache.set(self.node_key, nodes, self.CACHE_TIME)
        cache.set(self.system_key, system_users, self.CACHE_TIME)
        self.set_cache_meta()

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

    def expire_all_cache_meta(self):
        key = self.CACHE_META_KEY.format('*')
        cache.delete_pattern(key)

    @classmethod
    def expire_all_cache(cls):
        key = cls.CACHE_KEY.format('*', '*')
        cache.delete_pattern(key)


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


def parse_node_to_tree_node(node):
    from . import serializers
    name = '{} ({})'.format(node.value, node.assets_amount)
    node_serializer = serializers.GrantedNodeSerializer(node)
    data = {
        'id': node.key,
        'name': name,
        'title': name,
        'pId': node.parent_key,
        'isParent': True,
        'open': node.is_root(),
        'meta': {
            'node': node_serializer.data,
            'type': 'node'
        }
    }
    tree_node = TreeNode(**data)
    return tree_node


def parse_asset_to_tree_node(node, asset, system_users):
    system_users_protocol_matched = [s for s in system_users if s.protocol == asset.protocol]
    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_protocol_matched:
        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,
            'comment': system_user.comment,
        })
    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,
                'port': asset.port,
                'protocol': asset.protocol,
                'platform': asset.platform,
                'domain': None if not asset.domain else asset.domain.id,
                'is_active': asset.is_active,
                'comment': asset.comment
            },
        }
    }
    tree_node = TreeNode(**data)
    return tree_node