Perm cache (#3211)

* [Update] 修改用户树的cache

* [Update] 修改用户树缓存

* [Update] telnet (beta) => telnet

apps/assets/migrations/0038_auto_20190911_1634.py

@@ -0,0 +1,23 @@
+# Generated by Django 2.1.7 on 2019-09-11 08:34
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0037_auto_20190724_2002'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='asset',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'ssh'), ('rdp', 'rdp'), ('telnet', 'telnet'), ('vnc', 'vnc')], default='ssh', max_length=128, verbose_name='Protocol'),
+        ),
+        migrations.AlterField(
+            model_name='systemuser',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'ssh'), ('rdp', 'rdp'), ('telnet', 'telnet'), ('vnc', 'vnc')], default='ssh', max_length=16, verbose_name='Protocol'),
+        ),
+    ]

apps/assets/models/asset.py

@@ -57,7 +57,7 @@ class ProtocolsMixin:
     PROTOCOL_CHOICES = (
         (PROTOCOL_SSH, 'ssh'),
         (PROTOCOL_RDP, 'rdp'),
-        (PROTOCOL_TELNET, 'telnet (beta)'),
+        (PROTOCOL_TELNET, 'telnet'),
         (PROTOCOL_VNC, 'vnc'),
     )
 

apps/assets/models/node.py

@@ -11,7 +11,7 @@ from django.utils.translation import ugettext
 from django.core.cache import cache
 
 from orgs.mixins.models import OrgModelMixin, OrgManager
-from orgs.utils import set_current_org, get_current_org, tmp_to_root_org, tmp_to_org
+from orgs.utils import set_current_org, get_current_org, tmp_to_org
 from orgs.models import Organization
 
 
@@ -398,13 +398,24 @@ class Node(OrgModelMixin, SomeNodesMixin, TreeMixin, FamilyMixin, FullValueMixin
     def level(self):
         return len(self.key.split(':'))
 
+    @staticmethod
+    def refresh_user_tree_cache():
+        """
+        当节点-节点关系，节点-资产关系发生变化时，应该刷新用户授权树缓存
+        :return:
+        """
+        from perms.utils.asset_permission import AssetPermissionUtilV2
+        AssetPermissionUtilV2.expire_all_user_tree_cache()
+
     @classmethod
     def refresh_nodes(cls):
         cls.refresh_tree()
+        cls.refresh_user_tree_cache()
 
     @classmethod
     def refresh_assets(cls):
         cls.refresh_node_assets()
+        cls.refresh_user_tree_cache()
 
     def as_tree_node(self):
         from common.tree import TreeNode

apps/assets/models/user.py

@@ -96,7 +96,7 @@ class SystemUser(AssetUser):
     PROTOCOL_CHOICES = (
         (PROTOCOL_SSH, 'ssh'),
         (PROTOCOL_RDP, 'rdp'),
-        (PROTOCOL_TELNET, 'telnet (beta)'),
+        (PROTOCOL_TELNET, 'telnet'),
         (PROTOCOL_VNC, 'vnc'),
     )
 

apps/assets/signals_handler.py

@@ -108,7 +108,7 @@ def on_system_user_nodes_change(sender, instance=None, action=None, model=None,
     else:
         nodes_keys = [instance.key]
         system_users = queryset
-    assets = Node.get_nodes_all_assets(nodes_keys)
+    assets = Node.get_nodes_all_assets(nodes_keys).values_list('id', flat=True)
     for system_user in system_users:
         system_user.assets.add(*tuple(assets))
 
@@ -132,7 +132,7 @@ def on_asset_nodes_add(sender, instance=None, action='', model=None, pk_set=None
     if action != "post_add":
         return
     logger.debug("Assets node add signal recv: {}".format(action))
-    queryset = model.objects.filter(pk__in=pk_set)
+    queryset = model.objects.filter(pk__in=pk_set).values_list('id', flat=True)
     if model == Node:
         nodes = queryset
         assets = [instance]

apps/assets/templates/assets/_system_user.html

@@ -119,7 +119,7 @@ function autoLoginModeProtocol() {
         $(sudo_id).closest('.form-group').addClass('hidden');
         $(shell_id).closest('.form-group').addClass('hidden');
     }
-    else if (protocol === 'telnet (beta)') {
+    else if (protocol === 'telnet') {
         $('.auth-fields').removeClass('hidden');
         $(auto_generate_key).closest('.form-group').addClass('hidden');
         $(private_key_id).closest('.form-group').addClass('hidden');
@@ -165,7 +165,7 @@ function manualLoginModeProtocol() {
         $(sudo_id).closest('.form-group').addClass('hidden');
         $(shell_id).closest('.form-group').addClass('hidden');
     }
-    else if (protocol === 'telnet (beta)') {
+    else if (protocol === 'telnet') {
         $('.auth-fields').addClass('hidden');
         $(auto_generate_key).closest('.form-group').addClass('hidden');
         $(password_id).closest('.form-group').addClass('hidden');

apps/assets/utils.py

@@ -70,6 +70,7 @@ class TreeService(Tree):
     @classmethod
     @timeit
     def new(cls):
+        print("Call new")
         from .models import Node
         from orgs.utils import tmp_to_root_org
 

apps/common/signals_handlers.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 import re
+import os
 from collections import defaultdict
 from django.conf import settings
 from django.dispatch import receiver
@@ -12,8 +13,8 @@ from common.utils import get_logger
 from .local import thread_local
 
 pattern = re.compile(r'FROM `(\w+)`')
-# logger = logging.getLogger('jmsdb')
 logger = get_logger(__name__)
+DEBUG_DB_QUERY = os.environ.get('DEBUG_DB_QUERY', '0') == '1'
 
 
 class Counter:
@@ -57,7 +58,7 @@ def on_request_finished_release_local(sender, **kwargs):
     thread_local.__release_local__()
 
 
-if settings.DEBUG:
+if settings.DEBUG and DEBUG_DB_QUERY:
     request_finished.connect(on_request_finished_logging_db_query)
 
 

apps/perms/api/user_permission.py

@@ -45,13 +45,13 @@ class UserPermissionMixin:
 
     def initial(self, *args, **kwargs):
         super().initial(*args, *kwargs)
-        self.obj = self.get_object()
+        self.obj = self.get_obj()
 
     def get(self, request, *args, **kwargs):
         set_to_root_org()
         return super().get(request, *args, **kwargs)
 
-    def get_object(self):
+    def get_obj(self):
         user_id = self.kwargs.get('pk', '')
         if user_id:
             user = get_object_or_404(User, id=user_id)
@@ -65,12 +65,13 @@ class UserPermissionMixin:
         return super().get_permissions()
 
 
-class UserNodePermissionMixin(UserPermissionMixin):
+class UserAssetPermissionMixin(UserPermissionMixin):
     util = None
 
     def initial(self, *args, **kwargs):
         super().initial(*args, *kwargs)
-        self.util = AssetPermissionUtilV2(self.obj)
+        cache_policy = self.request.query_params.get('cache_policy', '0')
+        self.util = AssetPermissionUtilV2(self.obj, cache_policy=cache_policy)
 
 
 class UserNodeTreeMixin:
@@ -125,7 +126,7 @@ class UserAssetTreeMixin:
         return super().get_serializer(queryset, many=many, **kwargs)
 
 
-class UserGrantedAssetsApi(UserPermissionMixin, ListAPIView):
+class UserGrantedAssetsApi(UserAssetPermissionMixin, ListAPIView):
     permission_classes = (IsOrgAdminOrAppUser,)
     serializer_class = serializers.AssetGrantedSerializer
     only_fields = serializers.AssetGrantedSerializer.Meta.only_fields
@@ -151,8 +152,7 @@ class UserGrantedAssetsApi(UserPermissionMixin, ListAPIView):
         return queryset
 
     def get_queryset(self):
-        util = AssetPermissionUtilV2(self.obj)
-        queryset = util.get_assets().only(*self.only_fields)
+        queryset = self.util.get_assets().only(*self.only_fields)
         return queryset
 
 
@@ -165,13 +165,12 @@ class UserGrantedNodeAssetsApi(UserGrantedAssetsApi):
         node_id = self.kwargs.get("node_id")
         node = get_object_or_404(Node, pk=node_id)
         deep = self.request.query_params.get("all", "0") == "1"
-        util = AssetPermissionUtilV2(self.obj)
-        queryset = util.get_nodes_assets(node, deep=deep)\
+        queryset = self.util.get_nodes_assets(node, deep=deep)\
             .only(*self.only_fields)
         return queryset
 
 
-class UserGrantedNodesApi(UserNodePermissionMixin, ListAPIView):
+class UserGrantedNodesApi(UserAssetPermissionMixin, ListAPIView):
     """
     查询用户授权的所有节点的API
     """
@@ -267,55 +266,56 @@ class UserGrantedNodeChildrenWithAssetsAsTreeApi(UserGrantedNodeChildrenAsTreeAp
         return _queryset
 
 
-class GetUserAssetPermissionActionsApi(RetrieveAPIView):
+class GetUserAssetPermissionActionsApi(UserAssetPermissionMixin, RetrieveAPIView):
     permission_classes = (IsOrgAdminOrAppUser,)
     serializer_class = serializers.ActionsSerializer
 
-    def get_object(self):
+    def get_obj(self):
         user_id = self.request.query_params.get('user_id', '')
+        user = get_object_or_404(User, id=user_id)
+        return user
+
+    def get_object(self):
         asset_id = self.request.query_params.get('asset_id', '')
         system_id = self.request.query_params.get('system_user_id', '')
 
         try:
-            user_id = uuid.UUID(user_id)
             asset_id = uuid.UUID(asset_id)
             system_id = uuid.UUID(system_id)
         except ValueError:
             return Response({'msg': False}, status=403)
 
-        user = get_object_or_404(User, id=user_id)
         asset = get_object_or_404(Asset, id=asset_id)
         system_user = get_object_or_404(SystemUser, id=system_id)
 
-        util = AssetPermissionUtilV2(user)
-        system_users_actions = util.get_asset_system_users_with_actions(asset)
+        system_users_actions = self.util.get_asset_system_users_with_actions(asset)
         actions = system_users_actions.get(system_user)
         return {"actions": actions}
 
 
-class ValidateUserAssetPermissionApi(APIView):
+class ValidateUserAssetPermissionApi(UserAssetPermissionMixin, APIView):
     permission_classes = (IsOrgAdminOrAppUser,)
+
+    def get_obj(self):
+        user_id = self.request.query_params.get('user_id', '')
+        user = get_object_or_404(User, id=user_id)
+        return user
     
     def get(self, request, *args, **kwargs):
-        user_id = request.query_params.get('user_id', '')
         asset_id = request.query_params.get('asset_id', '')
         system_id = request.query_params.get('system_user_id', '')
         action_name = request.query_params.get('action_name', '')
-        cache_policy = self.request.query_params.get("cache_policy", '0')
 
         try:
-            user_id = uuid.UUID(user_id)
             asset_id = uuid.UUID(asset_id)
             system_id = uuid.UUID(system_id)
         except ValueError:
             return Response({'msg': False}, status=403)
 
-        user = get_object_or_404(User, id=user_id)
         asset = get_object_or_404(Asset, id=asset_id)
         system_user = get_object_or_404(SystemUser, id=system_id)
 
-        util = AssetPermissionUtilV2(user)
-        system_users_actions = util.get_asset_system_users_with_actions(asset)
+        system_users_actions = self.util.get_asset_system_users_with_actions(asset)
         actions = system_users_actions.get(system_user)
         if action_name in Action.value_to_choices(actions):
             return Response({'msg': True}, status=200)
@@ -326,21 +326,19 @@ class RefreshAssetPermissionCacheApi(RetrieveAPIView):
     permission_classes = (IsOrgAdmin,)
 
     def retrieve(self, request, *args, **kwargs):
-        # expire all cache
-        # AssetPermissionUtil.expire_all_cache()
+        AssetPermissionUtilV2.expire_all_user_tree_cache()
         return Response({'msg': True}, status=200)
 
 
-class UserGrantedAssetSystemUsersApi(UserPermissionMixin, ListAPIView):
+class UserGrantedAssetSystemUsersApi(UserAssetPermissionMixin, ListAPIView):
     permission_classes = (IsOrgAdminOrAppUser,)
     serializer_class = serializers.AssetSystemUserSerializer
     only_fields = serializers.AssetSystemUserSerializer.Meta.only_fields
 
     def get_queryset(self):
-        util = AssetPermissionUtilV2(self.obj)
         asset_id = self.kwargs.get('asset_id')
         asset = get_object_or_404(Asset, id=asset_id)
-        system_users_with_actions = util.get_asset_system_users_with_actions(asset)
+        system_users_with_actions = self.util.get_asset_system_users_with_actions(asset)
         system_users = []
         for system_user, actions in system_users_with_actions.items():
             system_user.actions = actions

apps/perms/signals_handler.py

@@ -6,25 +6,38 @@ from django.dispatch import receiver
 from common.utils import get_logger
 from common.decorator import on_transaction_commit
 from .models import AssetPermission
+from .utils.asset_permission import AssetPermissionUtilV2
 
 
 logger = get_logger(__file__)
 
+permission_m2m_senders = (
+    AssetPermission.nodes.through,
+    AssetPermission.assets.through,
+    AssetPermission.users.through,
+    AssetPermission.user_groups.through,
+)
+
 
-@receiver(post_save, sender=AssetPermission, dispatch_uid="my_unique_identifier")
 @on_transaction_commit
-def on_permission_created(sender, instance=None, created=False, **kwargs):
-    pass
+def on_permission_m2m_change(sender, action='', **kwargs):
+    if not action.startswith('post'):
+        return
+    logger.debug('Asset permission m2m changed, refresh user tree cache')
+    AssetPermissionUtilV2.expire_all_user_tree_cache()
 
 
-@receiver(post_save, sender=AssetPermission)
-def on_permission_update(sender, **kwargs):
-    pass
+for sender in permission_m2m_senders:
+    m2m_changed.connect(on_permission_m2m_change, sender=sender)
 
 
-@receiver(post_delete, sender=AssetPermission)
-def on_permission_delete(sender, **kwargs):
-    pass
+@receiver([post_save, post_delete], sender=AssetPermission)
+@on_transaction_commit
+def on_permission_change(sender, action='', **kwargs):
+    logger.debug('Asset permission changed, refresh user tree cache')
+    AssetPermissionUtilV2.expire_all_user_tree_cache()
+
+# Todo: 检查授权规则到期，从而修改授权规则
 
 
 @receiver(m2m_changed, sender=AssetPermission.nodes.through)
@@ -34,7 +47,7 @@ def on_permission_nodes_changed(sender, instance=None, action='', **kwargs):
     if isinstance(instance, AssetPermission):
         logger.debug("Asset permission nodes change signal received")
         nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
-        system_users = instance.system_users.all()
+        system_users = instance.system_users.all().values_list('id', flat=True)
         for system_user in system_users:
             system_user.nodes.add(*tuple(nodes))
 
@@ -46,7 +59,7 @@ def on_permission_assets_changed(sender, instance=None, action='', **kwargs):
     if isinstance(instance, AssetPermission):
         logger.debug("Asset permission assets change signal received")
         assets = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
-        system_users = instance.system_users.all()
+        system_users = instance.system_users.all().values_list('id', flat=True)
         for system_user in system_users:
             system_user.assets.add(*tuple(assets))
 
@@ -58,8 +71,8 @@ def on_permission_system_users_changed(sender, instance=None, action='', **kwarg
     if isinstance(instance, AssetPermission):
         system_users = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
         logger.debug("Asset permission system_users change signal received")
-        assets = instance.assets.all()
-        nodes = instance.nodes.all()
+        assets = instance.assets.all().values_list('id', flat=True)
+        nodes = instance.nodes.all().values_list('id', flat=True)
         for system_user in system_users:
             system_user.nodes.add(*tuple(nodes))
             system_user.assets.add(*tuple(assets))

apps/perms/utils/asset_permission.py

@@ -59,7 +59,55 @@ def get_system_user_permissions(system_user):
     )
 
 
-class AssetPermissionUtilV2:
+class AssetPermissionUtilCacheMixin:
+    user_tree_cache_key = 'USER_PERM_TREE_{}'
+    user_tree_cache_ttl = settings.ASSETS_PERM_CACHE_TIME
+    user_tree_cache_enable = settings.ASSETS_PERM_CACHE_ENABLE
+    cache_policy = '0'
+    obj_id = ''
+
+    def expire_user_tree_cache(self):
+        key = self.user_tree_cache_key.format(self.obj_id)
+        cache.delete(key)
+
+    @classmethod
+    def expire_all_user_tree_cache(cls):
+        key = cls.user_tree_cache_key.format('*')
+        cache.delete_pattern(key)
+
+    def set_user_tree_to_cache(self, user_tree):
+        data = pickle.dumps(user_tree)
+        key = self.user_tree_cache_key.format(self.obj_id)
+        cache.set(key, data, self.user_tree_cache_ttl)
+
+    def get_user_tree_from_cache(self):
+        key = self.user_tree_cache_key.format(self.obj_id)
+        data = cache.get(key)
+        if not data:
+            return None
+        user_tree = pickle.loads(data)
+        return user_tree
+
+    def get_user_tree_from_cache_if_need(self):
+        if not self.user_tree_cache_enable:
+            return None
+        if self.cache_policy == '1':
+            return self.get_user_tree_from_cache()
+        elif self.cache_policy == '2':
+            self.expire_user_tree_cache()
+            return None
+        else:
+            return None
+
+    def set_user_tree_to_cache_if_need(self, user_tree):
+        if self.cache_policy == '0':
+            return
+        if not self.user_tree_cache_enable:
+            return None
+        self.set_user_tree_to_cache(user_tree)
+
+
+class AssetPermissionUtilV2(AssetPermissionUtilCacheMixin):
     get_permissions_map = {
         "User": get_user_permissions,
         "UserGroup": get_user_group_permissions,
@@ -71,8 +119,6 @@ class AssetPermissionUtilV2:
         'id', 'hostname', 'ip', "platform", "domain_id",
         'comment', 'is_active', 'os', 'org_id'
     )
-    user_tree_cache_key = 'USER_PERM_TREE_{}'
-    user_tree_cache_ttl = 3600
 
     def __init__(self, obj, cache_policy='0'):
         self.object = obj
@@ -80,13 +126,8 @@ class AssetPermissionUtilV2:
         self.obj_id = str(obj.id)
         self._permissions = None
         self._permissions_id = None  # 标记_permission的唯一值
-        self._assets = None
         self._filter_id = 'None'  # 当通过filter更改 permission是标记
         self.change_org_if_need()
-        self.nodes = None
-        self._nodes = None
-        self._assets_direct = None
-        self._nodes_direct = None
         self._user_tree = None
         self.full_tree = Node.tree()
         self.mutex = threading.Lock()
@@ -109,31 +150,6 @@ class AssetPermissionUtilV2:
     def filter_permissions(self, **filters):
         self._permissions = self.permissions.filter(**filters)
 
-    @classmethod
-    def get_user_tree_from_cache(cls, obj_id):
-        return None
-        key = cls.user_tree_cache_key.format(obj_id)
-        data = cache.get(key)
-        if not data:
-            return None
-        user_tree = pickle.loads(data)
-        return user_tree
-
-    @classmethod
-    def expire_user_tree_cache(cls, obj_id):
-        if obj_id == 'all':
-            key = cls.user_tree_cache_key.format('*')
-            cache.delete_pattern(key)
-        else:
-            key = cls.user_tree_cache_key.format(obj_id)
-            cache.delete(key)
-
-    @classmethod
-    def set_user_tree_to_cache(cls, obj_id, user_tree):
-        data = pickle.dumps(user_tree)
-        key = cls.user_tree_cache_key.format(obj_id)
-        cache.set(key, data, cls.user_tree_cache_ttl)
-
     @property
     def user_tree(self):
         return self.get_user_tree()
@@ -268,11 +284,11 @@ class AssetPermissionUtilV2:
 
     @timeit
     def get_user_tree(self):
+        # 使用锁，保证多次获取tree的时候顺序执行，可以使用缓存
         with self.mutex:
             if self._user_tree:
                 return self._user_tree
-            print(id(self), self._user_tree)
-            user_tree = self.__class__.get_user_tree_from_cache(self.obj_id)
+            user_tree = self.get_user_tree_from_cache_if_need()
             if user_tree:
                 self._user_tree = user_tree
                 return user_tree
@@ -286,7 +302,8 @@ class AssetPermissionUtilV2:
             self.add_single_assets_node_to_user_tree(user_tree)
             self.parse_user_tree_to_full_tree(user_tree)
             self.add_empty_node_if_need(user_tree)
-            self.__class__.set_user_tree_to_cache(self.obj_id, user_tree)
+            self.set_user_tree_to_cache_if_need(user_tree)
+            self._user_tree = user_tree
             return user_tree
 
     # Todo: 是否可以获取多个资产的系统用户