mirror of https://github.com/jumpserver/jumpserver
418 lines
15 KiB
Python
418 lines
15 KiB
Python
import time
|
|
from collections import defaultdict
|
|
|
|
from django.conf import settings
|
|
from django.core.cache import cache
|
|
from django.db import transaction
|
|
from django.db.models import F
|
|
|
|
from assets.models import Asset
|
|
from assets.utils import NodeAssetsUtil
|
|
from common.decorators import merge_delay_run
|
|
from common.decorators import on_transaction_commit
|
|
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
|
|
)
|
|
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'
|
|
''.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
|
|
|
|
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):
|
|
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()
|
|
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(str(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 = {
|
|
str(node.id): 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(
|
|
str_asset_id=F('asset_id'),
|
|
str_node_id=F('node_id')
|
|
).values_list('str_asset_id', 'str_node_id')
|
|
asset_node_pairs = [
|
|
(str(asset_id), str(node_id))
|
|
for asset_id, node_id in 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
|