jumpserver/apps/perms/utils/user_perm_tree.py

414 lines
15 KiB
Python
Raw Normal View History

import time
from collections import defaultdict
from django.conf import settings
from django.core.cache import cache
from django.db import transaction
from assets.models import Asset
from assets.utils import NodeAssetsUtil
2023-02-08 02:14:09 +00:00
from common.db.models import output_as_string
from common.decorators import on_transaction_commit, merge_delay_run
2023-02-08 02:14:09 +00:00
from common.utils import get_logger
from common.utils.common import lazyproperty, timeit
from orgs.models import Organization
from orgs.utils import (
current_org,
tmp_to_org,
tmp_to_root_org
)
from perms.locks import UserGrantedTreeRebuildLock
from perms.models import (
AssetPermission,
UserAssetGrantedTreeNodeRelation,
PermNode
)
2023-02-08 02:14:09 +00:00
from users.models import User
from . import UserPermAssetUtil
from .permission import AssetPermissionUtil
logger = get_logger(__name__)
__all__ = [
'UserPermTreeRefreshUtil',
'UserPermTreeExpireUtil'
]
class _UserPermTreeCacheMixin:
""" 缓存数据 users: {org_id, org_id }, 记录用户授权树已经构建完成的组织集合 """
cache_key_template = 'perms.user.node_tree.built_orgs.user_id:{user_id}'
def get_cache_key(self, user_id):
return self.cache_key_template.format(user_id=user_id)
@lazyproperty
def client(self):
return cache.client.get_client(write=True)
class UserPermTreeRefreshUtil(_UserPermTreeCacheMixin):
""" 用户授权树刷新工具, 针对某一个用户授权树的刷新 """
def __init__(self, user):
self.user = user
@lazyproperty
def orgs(self):
return self.user.orgs.distinct()
@lazyproperty
def org_ids(self):
return [str(o.id) for o in self.orgs]
@lazyproperty
def cache_key_user(self):
return self.get_cache_key(self.user.id)
@lazyproperty
def cache_key_time(self):
key = 'perms.user.node_tree.built_time.{}'.format(self.user.id)
return key
@timeit
def refresh_if_need(self, force=False):
built_just_now = False if settings.ASSET_SIZE == 'small' else cache.get(self.cache_key_time)
if built_just_now:
logger.info('Refresh user perm tree just now, pass: {}'.format(built_just_now))
return
to_refresh_orgs = self.orgs if force else self._get_user_need_refresh_orgs()
if not to_refresh_orgs:
logger.info('Not have to refresh orgs')
return
logger.info("Delay refresh user orgs: {} {}".format(self.user, [o.name for o in to_refresh_orgs]))
sync = True if settings.ASSET_SIZE == 'small' else False
refresh_user_orgs_perm_tree.apply(sync=sync, user_orgs=((self.user, tuple(to_refresh_orgs)),))
with tmp_to_root_org():
refresh_user_favorite_assets.apply(sync=sync, users=(self.user,))
@timeit
def refresh_tree_manual(self):
"""
用来手动 debug
:return:
"""
built_just_now = cache.get(self.cache_key_time)
if built_just_now:
logger.info('Refresh just now, pass: {}'.format(built_just_now))
return
to_refresh_orgs = self._get_user_need_refresh_orgs()
if not to_refresh_orgs:
logger.info('Not have to refresh orgs for user: {}'.format(self.user))
return
self.perform_refresh_user_tree(to_refresh_orgs)
@timeit
def perform_refresh_user_tree(self, to_refresh_orgs):
# 再判断一次,毕竟构建树比较慢
built_just_now = cache.get(self.cache_key_time)
if built_just_now:
logger.info('Refresh user perm tree just now, pass: {}'.format(built_just_now))
return
self._clean_user_perm_tree_for_legacy_org()
if settings.ASSET_SIZE != 'small':
ttl = settings.PERM_TREE_REGEN_INTERVAL
cache.set(self.cache_key_time, int(time.time()), ttl)
lock = UserGrantedTreeRebuildLock(self.user.id)
got = lock.acquire(blocking=False)
if not got:
logger.info('User perm tree rebuild lock not acquired, pass')
return
try:
for org in to_refresh_orgs:
self._rebuild_user_perm_tree_for_org(org)
self._mark_user_orgs_refresh_finished(to_refresh_orgs)
finally:
lock.release()
def _rebuild_user_perm_tree_for_org(self, org):
with tmp_to_org(org):
start = time.time()
UserPermTreeBuildUtil(self.user).rebuild_user_perm_tree()
end = time.time()
logger.info(
'Refresh user perm tree: [{user}] org [{org}] {use_time:.2f}s'
2023-02-08 02:14:09 +00:00
''.format(user=self.user, org=org, use_time=end - start)
)
def _clean_user_perm_tree_for_legacy_org(self):
with tmp_to_root_org():
""" Clean user legacy org node relations """
user_relations = UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user)
user_legacy_org_relations = user_relations.exclude(org_id__in=self.org_ids)
user_legacy_org_relations.delete()
def _get_user_need_refresh_orgs(self):
cached_org_ids = self.client.smembers(self.cache_key_user)
cached_org_ids = {oid.decode() for oid in cached_org_ids}
to_refresh_org_ids = set(self.org_ids) - cached_org_ids
to_refresh_orgs = list(Organization.objects.filter(id__in=to_refresh_org_ids))
logger.info(f'Need to refresh orgs: {to_refresh_orgs}')
return to_refresh_orgs
2022-12-20 06:53:25 +00:00
def _mark_user_orgs_refresh_finished(self, orgs):
org_ids = [str(org.id) for org in orgs]
self.client.sadd(self.cache_key_user, *org_ids)
class UserPermTreeExpireUtil(_UserPermTreeCacheMixin):
""" 用户授权树过期工具 """
@lazyproperty
def cache_key_all_user(self):
return self.get_cache_key('*')
def expire_perm_tree_for_nodes_assets(self, node_ids, asset_ids):
node_perm_ids = AssetPermissionUtil().get_permissions_for_nodes(node_ids, flat=True)
asset_perm_ids = AssetPermissionUtil().get_permissions_for_assets(asset_ids, flat=True)
perm_ids = set(node_perm_ids) | set(asset_perm_ids)
self.expire_perm_tree_for_perms(perm_ids)
@tmp_to_root_org()
def expire_perm_tree_for_perms(self, perm_ids):
org_perm_ids = AssetPermission.objects.filter(id__in=perm_ids).values_list('org_id', 'id')
org_perms_mapper = defaultdict(set)
for org_id, perm_id in org_perm_ids:
org_perms_mapper[org_id].add(perm_id)
for org_id, perms_id in org_perms_mapper.items():
user_ids = AssetPermission.get_all_users_for_perms(perm_ids, flat=True)
self.expire_perm_tree_for_users_orgs(user_ids, [org_id])
def expire_perm_tree_for_user_group(self, user_group):
group_ids = [user_group.id]
org_ids = [user_group.org_id]
self.expire_perm_tree_for_user_groups_orgs(group_ids, org_ids)
def expire_perm_tree_for_user_groups_orgs(self, group_ids, org_ids):
user_ids = User.groups.through.objects \
.filter(usergroup_id__in=group_ids) \
.values_list('user_id', flat=True).distinct()
self.expire_perm_tree_for_users_orgs(user_ids, org_ids)
@on_transaction_commit
def expire_perm_tree_for_users_orgs(self, user_ids, org_ids):
2024-02-06 03:33:04 +00:00
user_ids = list(user_ids)
org_ids = [str(oid) for oid in org_ids]
with self.client.pipeline() as p:
for uid in user_ids:
cache_key = self.get_cache_key(uid)
p.srem(cache_key, *org_ids)
p.execute()
2024-02-04 09:21:09 +00:00
users_display = ','.join([str(i) for i in user_ids[:3]])
if len(user_ids) > 3:
users_display += '...'
orgs_display = ','.join([str(i) for i in org_ids[:3]])
if len(org_ids) > 3:
orgs_display += '...'
logger.info('Expire perm tree for users: [{}], orgs: [{}]'.format(users_display, orgs_display))
def expire_perm_tree_for_all_user(self):
keys = self.client.keys(self.cache_key_all_user)
with self.client.pipeline() as p:
for k in keys:
p.delete(k)
p.execute()
logger.info('Expire all user perm tree')
@merge_delay_run(ttl=20)
def refresh_user_orgs_perm_tree(user_orgs=()):
for user, orgs in user_orgs:
util = UserPermTreeRefreshUtil(user)
util.perform_refresh_user_tree(orgs)
@merge_delay_run(ttl=20)
def refresh_user_favorite_assets(users=()):
for user in users:
util = UserPermAssetUtil(user)
util.refresh_favorite_assets()
util.refresh_type_nodes_tree_cache()
class UserPermTreeBuildUtil(object):
node_only_fields = ('id', 'key', 'parent_key', 'org_id')
def __init__(self, user):
self.user = user
self.user_perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True)
# {key: node}
self._perm_nodes_key_node_mapper = {}
def rebuild_user_perm_tree(self):
with transaction.atomic():
self.clean_user_perm_tree()
if not self.user_perm_ids:
logger.info('User({}) not have permissions'.format(self.user))
return
self.compute_perm_nodes()
self.compute_perm_nodes_asset_amount()
self.create_mapping_nodes()
def clean_user_perm_tree(self):
UserAssetGrantedTreeNodeRelation.objects.filter(user=self.user).delete()
def compute_perm_nodes(self):
self._compute_perm_nodes_for_direct()
self._compute_perm_nodes_for_direct_asset_if_need()
self._compute_perm_nodes_for_ancestor()
def compute_perm_nodes_asset_amount(self):
""" 这里计算的是一个组织的授权树 """
computed = self._only_compute_root_node_assets_amount_if_need()
if computed:
return
nodekey_assetid_mapper = defaultdict(set)
org_id = current_org.id
for key in self.perm_node_keys_for_granted:
asset_ids = PermNode.get_all_asset_ids_by_node_key(org_id, key)
nodekey_assetid_mapper[key].update(asset_ids)
for asset_id, node_id in self.direct_asset_id_node_id_pairs:
node_key = self.perm_nodes_id_key_mapper.get(node_id)
if not node_key:
continue
nodekey_assetid_mapper[node_key].add(asset_id)
util = NodeAssetsUtil(self.perm_nodes, nodekey_assetid_mapper)
util.generate()
for node in self.perm_nodes:
assets_amount = util.get_assets_amount(node.key)
node.assets_amount = assets_amount
def create_mapping_nodes(self):
to_create = []
for node in self.perm_nodes:
relation = UserAssetGrantedTreeNodeRelation(
user=self.user,
node=node,
node_key=node.key,
node_parent_key=node.parent_key,
node_from=node.node_from,
node_assets_amount=node.assets_amount,
org_id=node.org_id
)
to_create.append(relation)
UserAssetGrantedTreeNodeRelation.objects.bulk_create(to_create)
def _compute_perm_nodes_for_direct(self):
""" 直接授权的节点(叶子节点)"""
for node in self.direct_nodes:
if self.has_any_ancestor_direct_permed(node):
continue
node.node_from = node.NodeFrom.granted
self._perm_nodes_key_node_mapper[node.key] = node
def _compute_perm_nodes_for_direct_asset_if_need(self):
""" 直接授权的资产所在的节点(叶子节点)"""
if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
return
for node in self.direct_asset_nodes:
if self.has_any_ancestor_direct_permed(node):
continue
if node.key in self._perm_nodes_key_node_mapper:
continue
node.node_from = node.NodeFrom.asset
self._perm_nodes_key_node_mapper[node.key] = node
def _compute_perm_nodes_for_ancestor(self):
""" 直接授权节点 和 直接授权资产所在节点 的所有祖先节点 (构造完整树) """
ancestor_keys = set()
for node in self._perm_nodes_key_node_mapper.values():
ancestor_keys.update(node.get_ancestor_keys())
ancestor_keys -= set(self._perm_nodes_key_node_mapper.keys())
ancestors = PermNode.objects.filter(key__in=ancestor_keys).only(*self.node_only_fields)
for node in ancestors:
node.node_from = node.NodeFrom.child
self._perm_nodes_key_node_mapper[node.key] = node
@lazyproperty
def perm_node_keys_for_granted(self):
keys = [
key for key, node in self._perm_nodes_key_node_mapper.items()
if node.node_from == node.NodeFrom.granted
]
return keys
@lazyproperty
def perm_nodes_id_key_mapper(self):
mapper = {
node.id.hex: node.key
for key, node in self._perm_nodes_key_node_mapper.items()
}
return mapper
def _only_compute_root_node_assets_amount_if_need(self):
if len(self.perm_nodes) != 1:
return False
root_node = self.perm_nodes[0]
if not root_node.is_org_root():
return False
if root_node.node_from != root_node.NodeFrom.granted:
return False
root_node.granted_assets_amount = len(root_node.get_all_asset_ids())
return True
@lazyproperty
def perm_nodes(self):
""" 授权树的所有节点 """
return list(self._perm_nodes_key_node_mapper.values())
def has_any_ancestor_direct_permed(self, node):
""" 任何一个祖先节点被直接授权 """
return bool(set(node.get_ancestor_keys()) & set(self.direct_node_keys))
@lazyproperty
def direct_node_keys(self):
""" 直接授权的节点 keys """
return {n.key for n in self.direct_nodes}
@lazyproperty
def direct_nodes(self):
""" 直接授权的节点 """
node_ids = AssetPermission.nodes.through.objects \
.filter(assetpermission_id__in=self.user_perm_ids) \
.values_list('node_id', flat=True).distinct()
nodes = PermNode.objects.filter(id__in=node_ids).only(*self.node_only_fields)
return nodes
@lazyproperty
def direct_asset_nodes(self):
""" 获取直接授权的资产所在的节点 """
node_ids = [node_id for asset_id, node_id in self.direct_asset_id_node_id_pairs]
nodes = PermNode.objects.filter(id__in=node_ids).distinct().only(*self.node_only_fields)
return nodes
@lazyproperty
def direct_asset_id_node_id_pairs(self):
""" 直接授权的资产 id 和 节点 id """
asset_node_pairs = Asset.nodes.through.objects \
.filter(asset_id__in=self.direct_asset_ids) \
.annotate(
2023-02-08 02:14:09 +00:00
str_asset_id=output_as_string('asset_id'),
str_node_id=output_as_string('node_id')
).values_list('str_asset_id', 'str_node_id')
asset_node_pairs = list(asset_node_pairs)
return asset_node_pairs
@lazyproperty
def direct_asset_ids(self):
""" 直接授权的资产 ids """
asset_ids = AssetPermission.assets.through.objects \
.filter(assetpermission_id__in=self.user_perm_ids) \
.values_list('asset_id', flat=True) \
.distinct()
return asset_ids