diff --git a/iast/base/agent.py b/iast/base/agent.py index 64f73108..104a8d31 100644 --- a/iast/base/agent.py +++ b/iast/base/agent.py @@ -65,6 +65,23 @@ def get_user_agent_pro(auth_users, bindId): return result +# add by song +def get_user_agent_pro_by_agent_id(ids): + agentInfo = IastAgent.objects.filter(id__in=ids).values("id", "bind_project_id", "server_id","server__container") + result = { + "pidArr": {}, + "serverArr": {}, + "serverNameArr": {} + } + if agentInfo: + for item in agentInfo: + + result["pidArr"][item['id']] = item['bind_project_id'] + result["serverNameArr"][item['id']] = item['server__container'] + result["serverArr"][item['id']] = item['server_id'] + return result + + def get_all_server(ids): alls = IastServer.objects.filter(id__in=ids).values("id", "container") result = {} diff --git a/iast/views/vuls.py b/iast/views/vuls.py index 5490d9bf..a914e9f4 100644 --- a/iast/views/vuls.py +++ b/iast/views/vuls.py @@ -3,7 +3,7 @@ # author:owefsad # software: PyCharm # project: lingzhi-webapi - +import time from dongtai.endpoint import R from dongtai.endpoint import UserEndPoint from dongtai.models.vul_level import IastVulLevel @@ -11,7 +11,7 @@ from dongtai.models.strategy import IastStrategyModel from iast.base.agent import get_agents_with_project, get_user_project_name, \ - get_user_agent_pro, get_all_server + get_user_agent_pro, get_all_server,get_user_agent_pro_by_agent_id from iast.base.project_version import get_project_version, get_project_version_by_id from iast.serializers.vul import VulSerializer from django.utils.translation import gettext_lazy as _ @@ -279,13 +279,10 @@ def get(self, request): q = ~Q(hook_type_id=0) queryset = queryset.filter(q) projects_info = get_user_project_name(auth_users) - agentArr = get_user_agent_pro(auth_users, projects_info.keys()) - agentPro = agentArr['pidArr'] - agentServer = agentArr['serverArr'] - server_ids = agentArr['server_ids'] - allServer = get_all_server(server_ids) + # server_ids = agentArr['server_ids'] allType = IastVulLevel.objects.all().order_by("id") allTypeArr = {} + if allType: for item in allType: allTypeArr[item.id] = item.name_value @@ -293,21 +290,31 @@ 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, page, page_size) + # need todo change datas = VulSerializer(page_data, many=True).data pro_length = len(datas) if pro_length > 0: + need_ids = [] + agent_ids = [] + for ind in range(pro_length): + item = datas[ind] + agent_ids.append(item['agent_id']) + # need_ids.append(agentServer.get(item['agent_id'], 0)) + # edit by song + agentArr = get_user_agent_pro_by_agent_id(agent_ids) + agentPro = agentArr['pidArr'] + # agentServer = agentArr['serverArr'] + agentServerName = agentArr['serverNameArr'] + # allServer = get_all_server(need_ids) + for index in range(pro_length): item = datas[index] item['index'] = index - item['project_name'] = projects_info.get( - agentPro.get(item['agent_id'], 0), - _("The application has not been binded")) + item['project_name'] = projects_info.get(agentPro.get(item['agent_id'], 0), _("The application has not been binded")) item['project_id'] = agentPro.get(item['agent_id'], 0) - item['server_name'] = allServer.get( - agentServer.get(item['agent_id'], 0), "JavaApplication") - item['server_type'] = VulSerializer.split_container_name( - item['server_name']) - + # item['server_name'] = allServer.get(agentServer.get(item['agent_id'], 0), "JavaApplication") + item['server_name'] = agentServerName.get(item['agent_id'], "JavaApplication") + item['server_type'] = VulSerializer.split_container_name( item['server_name']) item['level_type'] = item['level_id'] item['level'] = allTypeArr.get(item['level_id'], "") end['data'].append(item) diff --git a/webapi/settings.py b/webapi/settings.py index 33c7fd56..2c70bbef 100644 --- a/webapi/settings.py +++ b/webapi/settings.py @@ -119,6 +119,7 @@ def get_installed_apps(): USE_L10N = True MODELTRANSLATION_FALLBACK_LANGUAGES = ('zh', 'en') MIDDLEWARE = [ + # 'dongtai_agent_python.middlewares.django_middleware.FireMiddleware', 'django.middleware.gzip.GZipMiddleware', 'utils.CSPMiddleware', 'django.middleware.clickjacking.XFrameOptionsMiddleware',