[Update] 用户资产添加缓存

apps/jumpserver/conf.py

@@ -311,6 +311,7 @@ defaults = {
     'REDIS_PASSWORD': '',
     'REDIS_DB_CELERY': 3,
     'REDIS_DB_CACHE': 4,
+    'REDIS_DB_SESSION': 5,
     'CAPTCHA_TEST_MODE': None,
     'TOKEN_EXPIRATION': 3600 * 24,
     'DISPLAY_PER_PAGE': 25,
@@ -348,6 +349,7 @@ defaults = {
     'HTTP_BIND_HOST': '0.0.0.0',
     'HTTP_LISTEN_PORT': 8080,
     'LOGIN_LOG_KEEP_DAYS': 90,
+    'ASSETS_PERM_CACHE_TIME': 3600,
 }
 
 

apps/jumpserver/settings.py

@@ -141,6 +141,16 @@ SESSION_COOKIE_DOMAIN = CONFIG.SESSION_COOKIE_DOMAIN
 CSRF_COOKIE_DOMAIN = CONFIG.CSRF_COOKIE_DOMAIN
 SESSION_COOKIE_AGE = CONFIG.SESSION_COOKIE_AGE
 SESSION_EXPIRE_AT_BROWSER_CLOSE = CONFIG.SESSION_EXPIRE_AT_BROWSER_CLOSE
+SESSION_ENGINE = 'redis_sessions.session'
+SESSION_REDIS = {
+    'host': CONFIG.REDIS_HOST,
+    'port': CONFIG.REDIS_PORT,
+    'password': CONFIG.REDIS_PASSWORD,
+    'db': CONFIG.REDIS_DB_SESSION,
+    'prefix': 'auth_session',
+    'socket_timeout': 1,
+    'retry_on_timeout': False
+}
 
 MESSAGE_STORAGE = 'django.contrib.messages.storage.cookie.CookieStorage'
 # Database
@@ -559,3 +569,6 @@ SWAGGER_SETTINGS = {
 # Default email suffix
 EMAIL_SUFFIX = CONFIG.EMAIL_SUFFIX
 LOGIN_LOG_KEEP_DAYS = CONFIG.LOGIN_LOG_KEEP_DAYS
+
+# User or user group permission cache time, default 3600 seconds
+ASSETS_PERM_CACHE_TIME = CONFIG.ASSETS_PERM_CACHE_TIME

apps/perms/api/__init__.py

+# -*- coding: utf-8 -*-
+#
+
+from .permission import *
+from .user_permission import *
+from .user_group_permission import *

apps/perms/api/permission.py

+# -*- coding: utf-8 -*-
+#
+
+from django.utils import timezone
+from django.db.models import Q
+from rest_framework.views import Response
+from rest_framework.generics import RetrieveUpdateAPIView
+from rest_framework import viewsets
+from rest_framework.pagination import LimitOffsetPagination
+
+from common.permissions import IsOrgAdmin
+from common.utils import get_object_or_none
+from ..models import AssetPermission
+from ..hands import (
+    User, UserGroup, Asset, Node, SystemUser,
+)
+from .. import serializers
+
+
+__all__ = [
+    'AssetPermissionViewSet', 'AssetPermissionRemoveUserApi',
+    'AssetPermissionAddUserApi', 'AssetPermissionRemoveAssetApi',
+    'AssetPermissionAddAssetApi',
+]
+
+
+class AssetPermissionViewSet(viewsets.ModelViewSet):
+    """
+    资产授权列表的增删改查api
+    """
+    queryset = AssetPermission.objects.all()
+    serializer_class = serializers.AssetPermissionCreateUpdateSerializer
+    pagination_class = LimitOffsetPagination
+    filter_fields = ['name']
+    permission_classes = (IsOrgAdmin,)
+
+    def get_serializer_class(self):
+        if self.action in ("list", 'retrieve'):
+            return serializers.AssetPermissionListSerializer
+        return self.serializer_class
+
+    def filter_valid(self, queryset):
+        valid = self.request.query_params.get('is_valid', None)
+        if valid is None:
+            return queryset
+        if valid in ['0', 'N', 'false', 'False']:
+            valid = False
+        else:
+            valid = True
+        now = timezone.now()
+        if valid:
+            queryset = queryset.filter(is_active=True).filter(
+                date_start__lt=now, date_expired__gt=now,
+            )
+        else:
+            queryset = queryset.filter(
+                Q(is_active=False) |
+                Q(date_start__gt=now) |
+                Q(date_expired__lt=now)
+            )
+        return queryset
+
+    def filter_system_user(self, queryset):
+        system_user_id = self.request.query_params.get('system_user_id')
+        system_user_name = self.request.query_params.get('system_user')
+        if system_user_id:
+            system_user = get_object_or_none(SystemUser, pk=system_user_id)
+        elif system_user_name:
+            system_user = get_object_or_none(SystemUser, name=system_user_name)
+        else:
+            return queryset
+        if not system_user:
+            return queryset.none()
+        queryset = queryset.filter(system_users=system_user)
+        return queryset
+
+    def filter_node(self, queryset):
+        node_id = self.request.query_params.get('node_id')
+        node_name = self.request.query_params.get('node')
+        if node_id:
+            node = get_object_or_none(Node, pk=node_id)
+        elif node_name:
+            node = get_object_or_none(Node, name=node_name)
+        else:
+            return queryset
+        if not node:
+            return queryset.none()
+        nodes = node.get_ancestor(with_self=True)
+        queryset = queryset.filter(nodes__in=nodes)
+        return queryset
+
+    def filter_asset(self, queryset):
+        asset_id = self.request.query_params.get('asset_id')
+        hostname = self.request.query_params.get('hostname')
+        ip = self.request.query_params.get('ip')
+        if asset_id:
+            assets = Asset.objects.filter(pk=asset_id)
+        elif hostname:
+            assets = Asset.objects.filter(hostname=hostname)
+        elif ip:
+            assets = Asset.objects.filter(ip=ip)
+        else:
+            return queryset
+        if not assets:
+            return queryset.none()
+        inherit_nodes = set()
+        for asset in assets:
+            for node in asset.nodes.all():
+                inherit_nodes.update(set(node.get_ancestor(with_self=True)))
+        queryset = queryset.filter(Q(assets__in=assets) | Q(nodes__in=inherit_nodes))
+        return queryset
+
+    def filter_user(self, queryset):
+        user_id = self.request.query_params.get('user_id')
+        username = self.request.query_params.get('username')
+        if user_id:
+            user = get_object_or_none(User, pk=user_id)
+        elif username:
+            user = get_object_or_none(User, username=username)
+        else:
+            return queryset
+        if not user:
+            return queryset.none()
+
+    def filter_user_group(self, queryset):
+        user_group_id = self.request.query_params.get('user_group_id')
+        user_group_name = self.request.query_params.get('user_group')
+        if user_group_id:
+            group = get_object_or_none(UserGroup, pk=user_group_id)
+        elif user_group_name:
+            group = get_object_or_none(UserGroup, name=user_group_name)
+        else:
+            return queryset
+        if not group:
+            return queryset.none()
+        queryset = queryset.filter(user_groups=group)
+        return queryset
+
+    def filter_keyword(self, queryset):
+        keyword = self.request.query_params.get('search')
+        if not keyword:
+            return queryset
+        queryset = queryset.filter(name__icontains=keyword)
+        return queryset
+
+    def filter_queryset(self, queryset):
+        queryset = super().filter_queryset(queryset)
+        queryset = self.filter_valid(queryset)
+        queryset = self.filter_keyword(queryset)
+        queryset = self.filter_asset(queryset)
+        queryset = self.filter_node(queryset)
+        queryset = self.filter_system_user(queryset)
+        queryset = self.filter_user_group(queryset)
+        return queryset
+
+    def get_queryset(self):
+        return self.queryset.all()
+
+
+class AssetPermissionRemoveUserApi(RetrieveUpdateAPIView):
+    """
+    将用户从授权中移除，Detail页面会调用
+    """
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = serializers.AssetPermissionUpdateUserSerializer
+    queryset = AssetPermission.objects.all()
+
+    def update(self, request, *args, **kwargs):
+        perm = self.get_object()
+        serializer = self.serializer_class(data=request.data)
+        if serializer.is_valid():
+            users = serializer.validated_data.get('users')
+            if users:
+                perm.users.remove(*tuple(users))
+            return Response({"msg": "ok"})
+        else:
+            return Response({"error": serializer.errors})
+
+
+class AssetPermissionAddUserApi(RetrieveUpdateAPIView):
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = serializers.AssetPermissionUpdateUserSerializer
+    queryset = AssetPermission.objects.all()
+
+    def update(self, request, *args, **kwargs):
+        perm = self.get_object()
+        serializer = self.serializer_class(data=request.data)
+        if serializer.is_valid():
+            users = serializer.validated_data.get('users')
+            if users:
+                perm.users.add(*tuple(users))
+            return Response({"msg": "ok"})
+        else:
+            return Response({"error": serializer.errors})
+
+
+class AssetPermissionRemoveAssetApi(RetrieveUpdateAPIView):
+    """
+    将用户从授权中移除，Detail页面会调用
+    """
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = serializers.AssetPermissionUpdateAssetSerializer
+    queryset = AssetPermission.objects.all()
+
+    def update(self, request, *args, **kwargs):
+        perm = self.get_object()
+        serializer = self.serializer_class(data=request.data)
+        if serializer.is_valid():
+            assets = serializer.validated_data.get('assets')
+            if assets:
+                perm.assets.remove(*tuple(assets))
+            return Response({"msg": "ok"})
+        else:
+            return Response({"error": serializer.errors})
+
+
+class AssetPermissionAddAssetApi(RetrieveUpdateAPIView):
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = serializers.AssetPermissionUpdateAssetSerializer
+    queryset = AssetPermission.objects.all()
+
+    def update(self, request, *args, **kwargs):
+        perm = self.get_object()
+        serializer = self.serializer_class(data=request.data)
+        if serializer.is_valid():
+            assets = serializer.validated_data.get('assets')
+            if assets:
+                perm.assets.add(*tuple(assets))
+            return Response({"msg": "ok"})
+        else:
+            return Response({"error": serializer.errors})

apps/perms/api/user_group_permission.py

+# -*- coding: utf-8 -*-
+#
+
+from django.shortcuts import get_object_or_404
+from rest_framework.generics import (
+    ListAPIView, get_object_or_404,
+)
+
+from common.permissions import IsOrgAdmin, IsOrgAdminOrAppUser
+from common.tree import TreeNodeSerializer
+from orgs.utils import set_to_root_org
+from ..utils import (
+    AssetPermissionUtil, parse_asset_to_tree_node, parse_node_to_tree_node
+)
+from ..hands import (
+    AssetGrantedSerializer, UserGroup,  Node, NodeSerializer
+)
+from .. import serializers
+
+
+__all__ = [
+    'UserGroupGrantedAssetsApi', 'UserGroupGrantedNodesApi',
+    'UserGroupGrantedNodesWithAssetsApi', 'UserGroupGrantedNodeAssetsApi',
+    'UserGroupGrantedNodesWithAssetsAsTreeApi',
+]
+
+
+class UserGroupGrantedAssetsApi(ListAPIView):
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = AssetGrantedSerializer
+
+    def get_queryset(self):
+        user_group_id = self.kwargs.get('pk', '')
+        queryset = []
+
+        if not user_group_id:
+            return queryset
+
+        user_group = get_object_or_404(UserGroup, id=user_group_id)
+        util = AssetPermissionUtil(user_group)
+        assets = util.get_assets()
+        for k, v in assets.items():
+            k.system_users_granted = v
+            queryset.append(k)
+        return queryset
+
+
+class UserGroupGrantedNodesApi(ListAPIView):
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = NodeSerializer
+
+    def get_queryset(self):
+        group_id = self.kwargs.get('pk', '')
+        queryset = []
+
+        if group_id:
+            group = get_object_or_404(UserGroup, id=group_id)
+            util = AssetPermissionUtil(group)
+            nodes = util.get_nodes_with_assets()
+            return nodes.keys()
+        return queryset
+
+
+class UserGroupGrantedNodesWithAssetsApi(ListAPIView):
+    permission_classes = (IsOrgAdmin,)
+    serializer_class = serializers.NodeGrantedSerializer
+
+    def get_queryset(self):
+        user_group_id = self.kwargs.get('pk', '')
+        queryset = []
+
+        if not user_group_id:
+            return queryset
+
+        user_group = get_object_or_404(UserGroup, id=user_group_id)
+        util = AssetPermissionUtil(user_group)
+        nodes = util.get_nodes_with_assets()
+        for node, _assets in nodes.items():
+            assets = _assets.keys()
+            for asset, system_users in _assets.items():
+                asset.system_users_granted = system_users
+            node.assets_granted = assets
+            queryset.append(node)
+        return queryset
+
+
+class UserGroupGrantedNodesWithAssetsAsTreeApi(ListAPIView):
+    serializer_class = TreeNodeSerializer
+    permission_classes = (IsOrgAdminOrAppUser,)
+    show_assets = True
+    system_user_id = None
+
+    def change_org_if_need(self):
+        if self.request.user.is_superuser or \
+                self.request.user.is_app or \
+                self.kwargs.get('pk') is None:
+            set_to_root_org()
+
+    def get(self, request, *args, **kwargs):
+        self.show_assets = request.query_params.get('show_assets', '1') == '1'
+        self.system_user_id = request.query_params.get('system_user')
+        return super().get(request, *args, **kwargs)
+
+    def get_queryset(self):
+        self.change_org_if_need()
+        user_group_id = self.kwargs.get('pk', '')
+        queryset = []
+        group = get_object_or_404(UserGroup, id=user_group_id)
+        util = AssetPermissionUtil(group)
+        if self.system_user_id:
+            util.filter_permission_with_system_user(system_user=self.system_user_id)
+        nodes = util.get_nodes_with_assets()
+        for node, assets in nodes.items():
+            data = parse_node_to_tree_node(node)
+            queryset.append(data)
+            if not self.show_assets:
+                continue
+            for asset, system_users in assets.items():
+                data = parse_asset_to_tree_node(node, asset, system_users)
+                queryset.append(data)
+        queryset = sorted(queryset)
+        return queryset
+
+
+class UserGroupGrantedNodeAssetsApi(ListAPIView):
+    permission_classes = (IsOrgAdminOrAppUser,)
+    serializer_class = AssetGrantedSerializer
+
+    def get_queryset(self):
+        user_group_id = self.kwargs.get('pk', '')
+        node_id = self.kwargs.get('node_id')
+
+        user_group = get_object_or_404(UserGroup, id=user_group_id)
+        node = get_object_or_404(Node, id=node_id)
+        util = AssetPermissionUtil(user_group)
+        nodes = util.get_nodes_with_assets()
+        assets = nodes.get(node, [])
+        for asset, system_users in assets.items():
+            asset.system_users_granted = system_users
+        return assets

apps/perms/utils.py

@@ -3,6 +3,8 @@
 from __future__ import absolute_import, unicode_literals
 from collections import defaultdict
 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
@@ -97,10 +99,15 @@ class AssetPermissionUtil:
         "SystemUser": get_node_permissions,
     }
 
-    def __init__(self, obj):
+    CACHE_KEY = '_ASSET_PERM_CACHE_{}_{}'
+    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._permissions = None
         self._assets = None
+        self.cache_policy = cache_policy
 
     @property
     def permissions(self):
@@ -141,7 +148,7 @@ class AssetPermissionUtil:
                 )
         return assets
 
-    def get_assets(self):
+    def get_assets_without_cache(self):
         if self._assets:
             return self._assets
         assets = self.get_assets_direct()
@@ -155,7 +162,26 @@ class AssetPermissionUtil:
         self._assets = assets
         return self._assets
 
-    def get_nodes_with_assets(self):
+    def get_assets_from_cache(self):
+        cache_key = self.CACHE_KEY.format(str(self.object.id), 'ASSETS')
+        cached = cache.get(cache_key)
+        if cached:
+            return cached
+        assets = self.get_assets_without_cache()
+        self.expire_cache()
+        cache.set(cache_key, assets, self.CACHE_TIME)
+        return assets
+
+    def get_assets(self):
+        if 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:
+            return self.get_assets_without_cache()
+
+    def get_nodes_with_assets_without_cache(self):
         """
         返回节点并且包含资产
         {"node": {"assets": set("system_user")}}
@@ -167,13 +193,60 @@ class AssetPermissionUtil:
             tree.add_asset(asset, system_users)
         return tree.get_nodes()
 
-    def get_system_users(self):
+    def get_nodes_with_assets_from_cache(self):
+        cache_key = self.CACHE_KEY.format(str(self.object.id), 'NODES_WITH_ASSETS')
+        cached = cache.get(cache_key)
+        if cached:
+            return cached
+        nodes = self.get_nodes_with_assets_without_cache()
+        self.expire_cache()
+        cache.set(cache_key, nodes, self.CACHE_TIME)
+        return nodes
+
+    def get_nodes_with_assets(self):
+        if 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):
+        cache_key = self.CACHE_KEY.format(str(self.object.id), 'SYSTEM_USER')
+        cached = cache.get(cache_key)
+        if cached:
+            return cached
+        self.expire_cache()
+        system_users = self.get_system_user_without_cache()
+        cache.set(cache_key, system_users, self.CACHE_TIME)
+        return system_users
+
+    def get_system_users(self):
+        if 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()
+
+    def expire_cache(self):
+        cache_key = self.CACHE_KEY.format(str(self.object.id), '*')
+        cache.delete_pattern(cache_key)
+
+    @classmethod
+    def expire_all_cache(cls):
+        cache_key = cls.CACHE_KEY.format('*', '*')
+        cache.delete_pattern(cache_key)
+
 
 def is_obj_attr_has(obj, val, attrs=("hostname", "ip", "comment")):
     if not attrs:

apps/users/models/user.py

@@ -310,7 +310,8 @@ class User(AbstractUser):
         if not isinstance(remote_addr, bytes):
             remote_addr = remote_addr.encode("utf-8")
         remote_addr = base64.b16encode(remote_addr)  # .replace(b'=', '')
-        token = cache.get('%s_%s' % (self.id, remote_addr))
+        cache_key = '%s_%s' % (self.id, remote_addr)
+        token = cache.get(cache_key)
         if not token:
             token = uuid.uuid4().hex
         cache.set(token, self.id, expiration)

requirements/requirements.txt

@@ -80,3 +80,4 @@ python-ldap==3.1.0
 tencentcloud-sdk-python==3.0.40
 django-radius==1.3.3
 ipip-ipdb==1.2.1
+django-redis-sessions==0.6.1