Merge pull request #460 from HXSecurity/revert-450-optimise/vuln-list…

…-serilizer-query

Revert "optimise: vuln list serilizer query"

iast/serializers/vul.py

@@ -40,16 +40,26 @@ def split_container_name(name):
         return name
 
     def get_language(self, obj):
-        return obj['agent__language']
+        if obj['agent_id'] not in self.AGENT_LANGUAGE_MAP:
+            agent_model = IastAgent.objects.filter(id=obj['agent_id']).first()
+            if agent_model:
+                self.AGENT_LANGUAGE_MAP[obj['agent_id']] = agent_model.language
+        return self.AGENT_LANGUAGE_MAP[obj['agent_id']]
 
     def get_type(self, obj):
+        hook_type = HookType.objects.filter(pk=obj['hook_type_id']).first()
+        hook_type_name = hook_type.name if hook_type else None
+        strategy = IastStrategyModel.objects.filter(pk=obj['strategy_id']).first()
+        strategy_name = strategy.vul_name if strategy else None
         type_ = list(
-            filter(lambda x: x is not None, [obj['strategy__vul_name'], obj['hook_type__name']]))
+            filter(lambda x: x is not None, [strategy_name, hook_type_name]))
         return type_[0] if type_ else ''
 
     def get_status(self, obj):
-        status__name = obj.get('status__name',None)
-        return status__name if status__name else ''
+        status = IastVulnerabilityStatus.objects.filter(
+            pk=obj['status_id']).first()
+        return status.name if status else ''
+
 
 class VulForPluginSerializer(serializers.ModelSerializer):
     type = serializers.SerializerMethodField()

iast/views/details_id.py

@@ -126,13 +126,7 @@ def query(self, ids, request):
         auth_users = self.get_auth_users(request.user)
         auth_agents = self.get_auth_agents(auth_users)
         vuls = IastVulnerabilityModel.objects.filter(
-            pk__in=ids, agent__in=auth_agents).values(
-            'id', 'hook_type_id', 'url', 'uri', 'agent_id', 'level_id',
-            'http_method', 'top_stack', 'bottom_stack', 'taint_position',
-            'latest_time', 'first_time','strategy_id',
-            'status_id','strategy__vul_name','hook_type__name','status__name',
-            'agent__language'
-            ).all()
+            pk__in=ids, agent__in=auth_agents).values().all()
         return vuls
 
     @extend_schema_with_envcheck(

iast/views/vuls.py

@@ -209,7 +209,11 @@ def get(self, request):
         if language:
             auth_agents = auth_agents.filter(language=language)
 
-        queryset = IastVulnerabilityModel.objects.filter(agent__in=auth_agents)
+        queryset = IastVulnerabilityModel.objects.values(
+            'id', 'hook_type_id', 'url', 'uri', 'agent_id', 'level_id',
+            'http_method', 'top_stack', 'bottom_stack', 'taint_position',
+            'latest_time', 'first_time','strategy_id',
+            'status_id').filter(agent__in=auth_agents)
 
         level = request.query_params.get('level')
         if level:
@@ -288,15 +292,8 @@ def get(self, request):
 
         page = request.query_params.get('page', 1)
         page_size = request.query_params.get("pageSize", 20)
-        page_summary, page_data = self.get_paginator(queryset.only('id'), page, page_size)
-        vul_ids = [i.id for i in page_data]
-        datas = VulSerializer(IastVulnerabilityModel.objects.filter(pk__in=vul_ids).values(
-            'id', 'hook_type_id', 'url', 'uri', 'agent_id', 'level_id',
-            'http_method', 'top_stack', 'bottom_stack', 'taint_position',
-            'latest_time', 'first_time','strategy_id',
-            'status_id','strategy__vul_name','hook_type__name','status__name',
-            'agent__language'
-            ), many=True).data
+        page_summary, page_data = self.get_paginator(queryset, page, page_size)
+        datas = VulSerializer(page_data, many=True).data
         pro_length = len(datas)
         if pro_length > 0:
             for index in range(pro_length):