jumpserver/apps/perms/utils/asset/user_permission.py

from functools import reduce, wraps
from operator import or_, and_
from uuid import uuid4
import threading
import inspect

from django.conf import settings
from django.db.models import F, Q, Value, BooleanField
from django.utils.translation import ugettext_lazy as _

from common.http import is_true
from common.utils import get_logger
from common.const.distributed_lock_key import UPDATE_MAPPING_NODE_TASK_LOCK_KEY
from orgs.utils import tmp_to_root_org
from common.utils.timezone import dt_formater, now
from assets.models import Node, Asset, FavoriteAsset
from django.db.transaction import atomic
from orgs import lock
from perms.models import UserGrantedMappingNode, RebuildUserTreeTask, AssetPermission
from users.models import User

logger = get_logger(__name__)

ADD = 'add'
REMOVE = 'remove'

UNGROUPED_NODE_KEY = 'ungrouped'
UNGROUPED_NODE_VALUE = _('Ungrouped')
FAVORITE_NODE_KEY = 'favorite'
FAVORITE_NODE_VALUE = _('Favorite')

TMP_GRANTED_FIELD = '_granted'
TMP_ASSET_GRANTED_FIELD = '_asset_granted'
TMP_GRANTED_ASSETS_AMOUNT_FIELD = '_granted_assets_amount'


# 使用场景
# Asset.objects.filter(get_user_resources_q_granted_by_permissions(user))
def get_user_resources_q_granted_by_permissions(user: User):
    """
    获取用户关联的 asset permission 或者 用户组关联的 asset permission 获取规则,
    前提 AssetPermission 对象中的 related_name 为 granted_by_permissions
    :param user:
    :return:
    """
    _now = now()
    return reduce(and_, (
        Q(granted_by_permissions__date_start__lt=_now),
        Q(granted_by_permissions__date_expired__gt=_now),
        Q(granted_by_permissions__is_active=True),
        (
            Q(granted_by_permissions__users=user) |
            Q(granted_by_permissions__user_groups__users=user)
        )
    ))


# 使用场景
# `Node.objects.annotate(**node_annotate_mapping_node)`
node_annotate_mapping_node = {
    TMP_GRANTED_FIELD: F('mapping_nodes__granted'),
    TMP_ASSET_GRANTED_FIELD: F('mapping_nodes__asset_granted'),
    TMP_GRANTED_ASSETS_AMOUNT_FIELD: F('mapping_nodes__assets_amount')
}


# 使用场景
# `Node.objects.annotate(**node_annotate_set_granted)`
node_annotate_set_granted = {
    TMP_GRANTED_FIELD: Value(True, output_field=BooleanField()),
}


def is_direct_granted_by_annotate(node):
    return getattr(node, TMP_GRANTED_FIELD, False)


def is_asset_granted(node):
    return getattr(node, TMP_ASSET_GRANTED_FIELD, False)


def get_granted_assets_amount(node):
    return getattr(node, TMP_GRANTED_ASSETS_AMOUNT_FIELD, 0)


def set_granted(obj):
    setattr(obj, TMP_GRANTED_FIELD, True)


def set_asset_granted(obj):
    setattr(obj, TMP_ASSET_GRANTED_FIELD, True)


VALUE_TEMPLATE = '{stage}:{rand_str}:thread:{thread_name}:{thread_id}:{now}'


def _generate_value(stage=lock.DOING):
    cur_thread = threading.current_thread()

    return VALUE_TEMPLATE.format(
        stage=stage,
        thread_name=cur_thread.name,
        thread_id=cur_thread.ident,
        now=dt_formater(now()),
        rand_str=uuid4()
    )


def build_user_mapping_node_lock(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        call_args = inspect.getcallargs(func, *args, **kwargs)
        user = call_args.get('user')
        if user is None or not isinstance(user, User):
            raise ValueError('You function must have `user` argument')

        key = UPDATE_MAPPING_NODE_TASK_LOCK_KEY.format(user_id=user.id)
        doing_value = _generate_value()
        commiting_value = _generate_value(stage=lock.COMMITING)

        try:
            locked = lock.acquire(key, doing_value, timeout=600)
            if not locked:
                logger.error(f'update_mapping_node_task_locked_failed for user: {user.id}')
                raise lock.SomeoneIsDoingThis

            with atomic(savepoint=False):
                func(*args, **kwargs)
                ok = lock.change_lock_state_to_commiting(key, doing_value, commiting_value)
                if not ok:
                    logger.error(f'update_mapping_node_task_timeout for user: {user.id}')
                    raise lock.Timeout
        finally:
            lock.release(key, commiting_value, doing_value)
    return wrapper


@build_user_mapping_node_lock
def rebuild_user_mapping_nodes_if_need_with_lock(user: User):
    tasks = RebuildUserTreeTask.objects.filter(user=user)
    if tasks:
        tasks.delete()
        rebuild_user_mapping_nodes(user)


@build_user_mapping_node_lock
def rebuild_user_mapping_nodes_with_lock(user: User):
    rebuild_user_mapping_nodes(user)


@tmp_to_root_org()
def compute_tmp_mapping_node_from_perm(user: User):
    node_only_fields = ('id', 'key', 'parent_key', 'assets_amount')

    # 查询直接授权节点
    nodes = Node.objects.filter(
        get_user_resources_q_granted_by_permissions(user)
    ).distinct().only(*node_only_fields)
    granted_key_set = {_node.key for _node in nodes}

    def _has_ancestor_granted(node):
        """
        判断一个节点是否有授权过的祖先节点
        """
        ancestor_keys = set(node.get_ancestor_keys())
        return ancestor_keys & granted_key_set

    key2leaf_nodes_mapper = {}

    # 给授权节点设置 _granted 标识，同时去重
    for _node in nodes:
        if _has_ancestor_granted(_node):
            continue

        if _node.key not in key2leaf_nodes_mapper:
            set_granted(_node)
            key2leaf_nodes_mapper[_node.key] = _node

    # 查询授权资产关联的节点设置
    def process_direct_granted_assets():
        # 查询直接授权资产
        asset_ids = Asset.objects.filter(
            get_user_resources_q_granted_by_permissions(user)
        ).distinct().values_list('id', flat=True)
        # 查询授权资产关联的节点设置
        granted_asset_nodes = Node.objects.filter(
            assets__id__in=asset_ids
        ).distinct().only(*node_only_fields)

        # 给资产授权关联的节点设置 _asset_granted 标识，同时去重
        for _node in granted_asset_nodes:
            if _has_ancestor_granted(_node):
                continue

            if _node.key not in key2leaf_nodes_mapper:
                key2leaf_nodes_mapper[_node.key] = _node
            set_asset_granted(key2leaf_nodes_mapper[_node.key])

    if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
        process_direct_granted_assets()

    leaf_nodes = key2leaf_nodes_mapper.values()

    # 计算所有祖先节点
    ancestor_keys = set()
    for _node in leaf_nodes:
        ancestor_keys.update(_node.get_ancestor_keys())

    # 从祖先节点 key 中去掉同时也是叶子节点的 key
    ancestor_keys -= key2leaf_nodes_mapper.keys()
    # 查出祖先节点
    ancestors = Node.objects.filter(key__in=ancestor_keys).only(*node_only_fields)
    return [*leaf_nodes, *ancestors]


def create_mapping_nodes(user, nodes, clear=True):
    to_create = []
    for node in nodes:
        _granted = getattr(node, TMP_GRANTED_FIELD, False)
        _asset_granted = getattr(node, TMP_ASSET_GRANTED_FIELD, False)
        _granted_assets_amount = getattr(node, TMP_GRANTED_ASSETS_AMOUNT_FIELD, 0)
        to_create.append(UserGrantedMappingNode(
            user=user,
            node=node,
            key=node.key,
            parent_key=node.parent_key,
            granted=_granted,
            asset_granted=_asset_granted,
            assets_amount=_granted_assets_amount,
        ))

    if clear:
        UserGrantedMappingNode.objects.filter(user=user).delete()
    UserGrantedMappingNode.objects.bulk_create(to_create)


def set_node_granted_assets_amount(user, node):
    """
    不依赖`UserGrantedMappingNode`直接查询授权计算资产数量
    """
    _granted = getattr(node, TMP_GRANTED_FIELD, False)
    if _granted:
        assets_amount = node.assets_amount
    else:
        if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
            assets_amount = count_direct_granted_node_assets(user, node.key)
        else:
            assets_amount = count_node_all_granted_assets(user, node.key)
    setattr(node, TMP_GRANTED_ASSETS_AMOUNT_FIELD, assets_amount)


def rebuild_user_mapping_nodes(user):
    logger.info(f'>>> {dt_formater(now())} start rebuild {user} mapping nodes')
    tmp_nodes = compute_tmp_mapping_node_from_perm(user)
    for _node in tmp_nodes:
        set_node_granted_assets_amount(user, _node)
    create_mapping_nodes(user, tmp_nodes)
    logger.info(f'>>> {dt_formater(now())} end rebuild {user} mapping nodes')


def rebuild_all_user_mapping_nodes():
    from users.models import User
    users = User.objects.all()
    for user in users:
        rebuild_user_mapping_nodes(user)


def get_user_granted_nodes_list_via_mapping_node(user):
    """
    这里的 granted nodes, 是整棵树需要的node，推算出来的也算
    :param user:
    :return:
    """
    # 获取 `UserGrantedMappingNode` 中对应的 `Node`
    nodes = Node.objects.filter(
        mapping_nodes__user=user,
    ).annotate(
        **node_annotate_mapping_node
    ).distinct()

    key_to_node_mapper = {}
    nodes_descendant_q = Q()

    for node in nodes:
        if not is_direct_granted_by_annotate(node):
            # 未授权的节点资产数量设置为 `UserGrantedMappingNode` 中的数量
            node.assets_amount = get_granted_assets_amount(node)
        else:
            # 直接授权的节点
            # 增加查询后代节点的过滤条件
            nodes_descendant_q |= Q(key__startswith=f'{node.key}:')
        key_to_node_mapper[node.key] = node

    if nodes_descendant_q:
        descendant_nodes = Node.objects.filter(
            nodes_descendant_q
        ).annotate(
            **node_annotate_set_granted
        )
        for node in descendant_nodes:
            key_to_node_mapper[node.key] = node

    all_nodes = key_to_node_mapper.values()
    return all_nodes


def get_user_granted_all_assets(user, via_mapping_node=True):
    asset_perms_id = get_user_all_assetpermissions_id(user)
    if via_mapping_node:
        granted_node_keys = UserGrantedMappingNode.objects.filter(
            user=user, granted=True,
        ).values_list('key', flat=True).distinct()
    else:
        granted_node_keys = Node.objects.filter(
            granted_by_permissions__id__in=asset_perms_id
        ).distinct().values_list('key', flat=True)
    granted_node_keys = Node.clean_children_keys(granted_node_keys)

    granted_node_q = Q()
    for _key in granted_node_keys:
        granted_node_q |= Q(nodes__key__startswith=f'{_key}:')
        granted_node_q |= Q(nodes__key=_key)

    assets__id = get_user_direct_granted_assets(user, asset_perms_id).values_list('id', flat=True)

    q = granted_node_q | Q(id__in=list(assets__id))
    return Asset.org_objects.filter(q).distinct()


def get_node_all_granted_assets(user: User, key):
    """
    此算法依据 `UserGrantedMappingNode` 的数据查询
    1. 查询该节点下的直接授权节点
    2. 查询该节点下授权资产关联的节点
    """

    assets = Asset.objects.none()

    # 查询该节点下的授权节点
    granted_mapping_nodes = UserGrantedMappingNode.objects.filter(
        user=user, granted=True,
    ).filter(
        Q(key__startswith=f'{key}:') | Q(key=key)
    )

    # 根据授权节点构建资产查询条件
    granted_nodes_qs = []
    for _node in granted_mapping_nodes:
        granted_nodes_qs.append(Q(nodes__key__startswith=f'{_node.key}:'))
        granted_nodes_qs.append(Q(nodes__key=_node.key))

    # 查询该节点下的资产授权节点
    only_asset_granted_mapping_nodes = UserGrantedMappingNode.objects.filter(
        user=user,
        asset_granted=True,
        granted=False,
    ).filter(Q(key__startswith=f'{key}:') | Q(key=key))

    # 根据资产授权节点构建查询
    only_asset_granted_nodes_qs = []
    for _node in only_asset_granted_mapping_nodes:
        only_asset_granted_nodes_qs.append(Q(nodes__id=_node.node_id))

    q = []
    if granted_nodes_qs:
        q.append(reduce(or_, granted_nodes_qs))

    if only_asset_granted_nodes_qs:
        only_asset_granted_nodes_q = reduce(or_, only_asset_granted_nodes_qs)
        only_asset_granted_nodes_q &= get_user_resources_q_granted_by_permissions(user)
        q.append(only_asset_granted_nodes_q)

    if q:
        assets = Asset.objects.filter(reduce(or_, q)).distinct()
    return assets


def get_direct_granted_node_ids(user: User, key):
    granted_q = get_user_resources_q_granted_by_permissions(user)

    # 先查出该节点下的直接授权节点
    granted_nodes = Node.objects.filter(
        Q(key__startswith=f'{key}:') | Q(key=key)
    ).filter(granted_q).distinct().only('id', 'key')

    node_ids = set()
    # 根据直接授权节点查询他们的子节点
    q = Q()
    for _node in granted_nodes:
        q |= Q(key__startswith=f'{_node.key}:')
        node_ids.add(_node.id)

    if q:
        descendant_ids = Node.objects.filter(q).values_list('id', flat=True).distinct()
        node_ids.update(descendant_ids)
    return node_ids


def get_node_all_granted_assets_from_perm(user: User, key):
    """
    此算法依据 `AssetPermission` 的数据查询
    1. 查询该节点下的直接授权节点
    2. 查询该节点下授权资产关联的节点
    """
    granted_q = get_user_resources_q_granted_by_permissions(user)
    # 直接授权资产查询条件
    q = (Q(nodes__key__startswith=f'{key}:') | Q(nodes__key=key)) & granted_q
    node_ids = get_direct_granted_node_ids(user, key)
    q |= Q(nodes__id__in=node_ids)
    asset_qs = Asset.objects.filter(q).distinct()
    return asset_qs


def get_direct_granted_node_assets_from_perm(user: User, key):
    node_ids = get_direct_granted_node_ids(user, key)
    asset_qs = Asset.objects.filter(nodes__id__in=node_ids).distinct()
    return asset_qs


def count_node_all_granted_assets(user: User, key):
    return get_node_all_granted_assets_from_perm(user, key).count()


def count_direct_granted_node_assets(user: User, key):
    return get_direct_granted_node_assets_from_perm(user, key).count()


def get_indirect_granted_node_children(user, key=''):
    """
    获取用户授权树中未授权节点的子节点
    只匹配在 `UserGrantedMappingNode` 中存在的节点
    """
    nodes = Node.objects.filter(
        mapping_nodes__user=user,
        parent_key=key
    ).annotate(
        _granted_assets_amount=F('mapping_nodes__assets_amount'),
        _granted=F('mapping_nodes__granted')
    ).distinct()

    # 设置节点授权资产数量
    for _node in nodes:
        if not is_direct_granted_by_annotate(_node):
            _node.assets_amount = get_granted_assets_amount(_node)
    return nodes


def get_top_level_granted_nodes(user):
    nodes = list(get_indirect_granted_node_children(user, key=''))
    if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
        ungrouped_node = get_ungrouped_node(user)
        nodes.insert(0, ungrouped_node)
    favorite_node = get_favorite_node(user)
    nodes.insert(0, favorite_node)
    return nodes


def get_user_all_assetpermissions_id(user: User):
    asset_perms_id = AssetPermission.objects.valid().filter(
        Q(users=user) | Q(user_groups__users=user)
    ).distinct().values_list('id', flat=True)
    return asset_perms_id


def get_user_direct_granted_assets(user, asset_perms_id=None):
    if asset_perms_id is None:
        asset_perms_id = get_user_all_assetpermissions_id(user)
    assets = Asset.org_objects.filter(granted_by_permissions__id__in=asset_perms_id).distinct()
    return assets


def count_user_direct_granted_assets(user):
    count = get_user_direct_granted_assets(user).values_list('id').count()
    return count


def get_ungrouped_node(user):
    assets_amount = count_user_direct_granted_assets(user)
    return Node(
        id=UNGROUPED_NODE_KEY,
        key=UNGROUPED_NODE_KEY,
        value=UNGROUPED_NODE_VALUE,
        assets_amount=assets_amount
    )


def get_favorite_node(user):
    assets_amount = FavoriteAsset.get_user_favorite_assets(user)\
        .values_list('id')\
        .count()
    return Node(
        id=FAVORITE_NODE_KEY,
        key=FAVORITE_NODE_KEY,
        value=FAVORITE_NODE_VALUE,
        assets_amount=assets_amount
    )


def rebuild_user_tree_if_need(request, user):
    """
    升级授权树策略后，用户的数据可能还未初始化，为防止用户显示没有数据
    先检查 MappingNode 如果没有数据，同步创建用户授权树
    """
    if is_true(request.query_params.get('rebuild_tree')) or \
            not UserGrantedMappingNode.objects.filter(user=user).exists():
        try:
            rebuild_user_mapping_nodes_with_lock(user)
        except lock.SomeoneIsDoingThis:
            # 您的数据正在初始化，请稍等
            raise lock.SomeoneIsDoingThis(
                detail=_('Please wait while your data is being initialized'),
                code='rebuild_tree_conflict'
            )