From 600b5e79a34ce2f7070ac84bc41bd1888ae49486 Mon Sep 17 00:00:00 2001
From: "xianzelin@huoxian.cn" <xianzelin@huoxian.cn>
Date: Fri, 7 Jan 2022 11:22:15 +0800
Subject: [PATCH] Enhancement: optimise agent_id query logic

---
 iast/base/agent.py                  | 18 +++++++++---------
 iast/views/project_report_export.py | 15 +++++++++------
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/iast/base/agent.py b/iast/base/agent.py
index dc721626..edc131c4 100644
--- a/iast/base/agent.py
+++ b/iast/base/agent.py
@@ -27,15 +27,15 @@ def get_agents_with_project(project_name, users):
     """
     agent_ids = []
     if project_name and project_name != '':
-        project_queryset = IastProject.objects.filter(user__in=users, name__icontains=project_name).values("id")
-        project_ids = []
+        project_ids = IastProject.objects.filter(
+            user__in=users,
+            name__icontains=project_name).values_list("id", flat=True).all()
 
-        if project_queryset:
-            for pro_item in project_queryset:
-                project_ids.append(pro_item['id'])
+        if project_ids:
 
-            relations = IastAgent.objects.filter(bind_project_id__in=project_ids).values("id")
-            agent_ids = [relation['id'] for relation in relations]
+            agent_ids = IastAgent.objects.filter(
+                bind_project_id__in=project_ids).values_list("id",
+                                                             flat=True).all()
 
     return agent_ids
 
@@ -143,9 +143,9 @@ def get_vul_count_by_agent(agent_ids, vid, user):
         typeArr = {}
         typeLevel = {}
         for one in typeInfo:
-            hook_type = hooktypes.get('hook_type_id', None)
+            hook_type = hooktypes.get(one['hook_type_id'], None)
             hook_type_name = hook_type['name'] if hook_type else None
-            strategy = strategys.get('strategy_id', None)
+            strategy = strategys.get(one['strategy_id'], None)
             strategy_name = strategy['vul_name'] if strategy else None
             type_ = list(
                 filter(lambda x: x is not None, [strategy_name, hook_type_name]))
diff --git a/iast/views/project_report_export.py b/iast/views/project_report_export.py
index b694d3b6..d207e3e3 100644
--- a/iast/views/project_report_export.py
+++ b/iast/views/project_report_export.py
@@ -26,10 +26,13 @@
 
 class _ProjectReportExportQuerySerializer(serializers.Serializer):
     vid = serializers.CharField(
-        help_text=_("The version id of the project"))
+        help_text=_("The version id of the project"),
+        required=False)
     pname = serializers.CharField(
-        help_text=_("The name of the project"))
-    pid = serializers.IntegerField(help_text=_("The id of the project"))
+        help_text=_("The name of the project"),
+        required=False)
+    pid = serializers.IntegerField(help_text=_("The id of the project"),
+        required=False)
 
 class ProjectReportExport(UserEndPoint):
     name = 'api-v1-word-maker'
@@ -46,9 +49,9 @@ def get_agents_with_project_id(pid, auth_users):
         :param auth_users:
         :return:
         """
-        relations = IastAgent.objects.filter(bind_project_id=pid,
-                                             user__in=auth_users).values("id")
-        agent_ids = [relation['id'] for relation in relations]
+        agent_ids = IastAgent.objects.filter(bind_project_id=pid,
+                                             user__in=auth_users).values_list(
+                                                 "id", flat=True).all()
         return agent_ids
 
     @staticmethod