From 40f103f995b32ffe25799f758fd0eb349a052311 Mon Sep 17 00:00:00 2001
From: "xianzelin@huoxian.cn" <xianzelin@huoxian.cn>
Date: Thu, 2 Dec 2021 19:45:06 +0800
Subject: [PATCH] fix permissions for sensitive information

---
 iast/views/sensitive_info_rule.py | 47 ++++++++++++++++++-------------
 1 file changed, 28 insertions(+), 19 deletions(-)

diff --git a/iast/views/sensitive_info_rule.py b/iast/views/sensitive_info_rule.py
index b1f04313..f50e6ad4 100644
--- a/iast/views/sensitive_info_rule.py
+++ b/iast/views/sensitive_info_rule.py
@@ -3,7 +3,7 @@
 # @file        : sensitive_info_rule
 # @created     : 星期三 11月 17, 2021 16:15:57 CST
 #
-# @description : 
+# @description :
 ######################################################################
 
 
@@ -42,13 +42,13 @@ class SensitiveInfoRuleSerializer(serializers.ModelSerializer):
     class Meta:
         model = IastSensitiveInfoRule
         fields = ['id', 'strategy_name','strategy_id','pattern_type_id','pattern_type_name','pattern','status','latest_time']
-    
+
     def get_strategy_name(self,obj):
         return obj.strategy.vul_name
 
     def get_strategy_id(self,obj):
         return obj.strategy.id
-    
+
     def get_pattern_type_id(self,obj):
         return obj.pattern_type.id
     def get_pattern_type_name(self,obj):
@@ -60,9 +60,9 @@ class Meta:
         fields = ['id', 'name', 'url']
 
     def get_url(self,obj):
-        url_dict = {1:'regex',2:'json'} 
+        url_dict = {1:'regex',2:'json'}
         return url_dict.get(obj.id,'')
-    
+
 
 class SensitiveInfoRuleCreateSerializer(serializers.Serializer):
     strategy_id = serializers.IntegerField(required=True)
@@ -83,16 +83,16 @@ class _SensitiveInfoArgsSerializer(serializers.Serializer):
 class _RegexPatternValidationSerializer(serializers.Serializer):
     pattern = serializers.CharField(help_text=_('regex pattern'))
     test_data = serializers.CharField(help_text=_('the data for test regex'))
-  
+
 class SensitiveInfoRuleViewSet(UserEndPoint,viewsets.ViewSet):
-    
+
     permission_classes_by_action = {'destory':(TalentAdminPermission,),}
 
     def get_permissions(self):
-      try:
-        return [permission() for permission in self.permission_classes_by_action[self.action]]
-      except KeyError:
-        return [permission() for permission in self.permission_classes]
+        try:
+            return [permission() for permission in self.permission_classes_by_action[self.action]]
+        except KeyError:
+            return [permission() for permission in self.permission_classes]
 
     @extend_schema_with_envcheck(
         [_SensitiveInfoArgsSerializer],
@@ -102,7 +102,7 @@ def get_permissions(self):
         _("Get the item corresponding to the user, support fuzzy search based on name."
           ),
     )
-    def list(self,request): 
+    def list(self,request):
         ser = _SensitiveInfoArgsSerializer(data=request.data)
         try:
             if ser.is_valid(True):
@@ -121,7 +121,7 @@ def list(self,request):
             queryset = queryset.filter(name__icontains=name)
         page_summary, page_data = self.get_paginator(queryset, page, page_size)
         return R.success(data=SensitiveInfoRuleSerializer(page_data,many=True).data,page=page_summary)
-    
+
     @extend_schema_with_envcheck(
             request=SensitiveInfoRuleCreateSerializer,
         tags=[_('SensitiveInfoRule')],
@@ -152,7 +152,7 @@ def create(self,request):
                     pattern_type=pattern_type,
                     pattern=pattern,
                     status=status,
-                    user=request.user) 
+                    user=request.user)
             return R.success(msg='create success',data=SensitiveInfoRuleSerializer(obj).data)
         else:
             return R.failure()
@@ -174,8 +174,12 @@ def update(self, request, pk):
                 status = ser.validated_data['status']
         except ValidationError as e:
             return R.failure(data=e.detail)
-        obj = IastSensitiveInfoRule.objects.filter(pk=pk).update(**ser.validated_data,latest_time=time.time())
+        users = self.get_auth_users(request.user)
+        obj = IastSensitiveInfoRule.objects.filter(
+            pk=pk, user__in=users).update(**ser.validated_data,
+                                          latest_time=time.time())
         return R.success(msg='update success')
+
     @extend_schema_with_envcheck(
         tags=[_('SensitiveInfoRule')],
         summary=_('SensitiveInfoRule delete'),
@@ -184,7 +188,9 @@ def update(self, request, pk):
           ),
     )
     def destory(self, request, pk):
-        IastSensitiveInfoRule.objects.filter(pk=pk).update(status=-1)
+        users = self.get_auth_users(request.user)
+        IastSensitiveInfoRule.objects.filter(pk=pk,
+                                             user__in=users).update(status=-1)
         return R.success(msg='delete success')
 
     @extend_schema_with_envcheck(
@@ -195,8 +201,11 @@ def destory(self, request, pk):
           ),
     )
     def retrieve(self, request, pk):
-        obj = IastSensitiveInfoRule.objects.filter(pk=pk,user=request.user).first()
+        users = self.get_auth_users(request.user)
+        obj = IastSensitiveInfoRule.objects.filter(pk=pk, user=users).first()
         return R.success(data=SensitiveInfoRuleSerializer(obj).data)
+
+
 class SensitiveInfoPatternTypeView(UserEndPoint):
 
     @extend_schema_with_envcheck(
@@ -257,7 +266,7 @@ def regextest(test_data,pattern):
         print(e)
         data = ''
         status = 0
-        return data,status 
+        return data,status
     ret = regex.findall(test_data)
     data = ret[0] if ret else ['']
     return data,1
@@ -269,4 +278,4 @@ def jsontest(test_data,pattern):
         print(e)
         data = ''
         status = 0
-    return data, status 
+    return data, status