diff --git a/iast/base/agent.py b/iast/base/agent.py index edc131c4..10be53c1 100644 --- a/iast/base/agent.py +++ b/iast/base/agent.py @@ -7,6 +7,7 @@ import json import re import time +from django.db.models import Count from dongtai.models.agent import IastAgent from dongtai.models.project import IastProject @@ -17,7 +18,7 @@ from iast.base.project_version import get_project_version from dongtai.models.strategy import IastStrategyModel from dongtai.models.project_version import IastProjectVersion - +from functools import reduce def get_agents_with_project(project_name, users): """ @@ -88,23 +89,35 @@ def get_project_vul_count(users, queryset, auth_agents, project_id=None): current_version=1, user__in=users).values_list('id', 'project_id').all() versions_map = {version[1]: version[0] for version in versions} + qss = [] + project_pair = {} for project in project_queryset: project_id = project['id'] version_id = versions_map.get(project_id, 0) agent_queryset = auth_agents.filter(project_version_id=version_id, bind_project_id=project_id) - count = queryset.filter(agent__in=agent_queryset).values('id').count() - if count is False: - result.append({ - "project_name": project['name'], - "count": 0, - "id": project_id - }) - else: + qs = queryset.filter(agent__in=agent_queryset).extra({ + "project_id": + project_id, + }).annotate(count=Count('id')).values('project_id', + 'count') + project_pair[project_id] = project['name'] + qss.append(qs) + if len(qss) > 1: + unionqs = qss[0].union(*qss[1::]) + elif len(qss) == 1: + unionqs = qss[0] + else: + unionqs = None + if unionqs: + for project_result in unionqs: result.append({ - "project_name": project['name'], - "count": count, - "id": project_id + "project_name": + project_pair[project_result['project_id']], + "count": + project_result['count'], + "project_id": + project_result['project_id'], }) result = sorted(result, key=lambda item: item['count'], reverse=True)[:5] return result diff --git a/iast/views/vul_summary.py b/iast/views/vul_summary.py index aed3fef5..f0d442aa 100644 --- a/iast/views/vul_summary.py +++ b/iast/views/vul_summary.py @@ -298,7 +298,7 @@ def get(self, request): vul_type = request.query_params.get('type') if vul_type: hook_types = HookType.objects.filter(name=vul_type).all() - strategys = IastStrategyModel.objects.filter(vul_name=vul_type).all() + strategys = IastStrategyModel.objects.filter(vul_name=vul_type).all() q = Q(hook_type__in=hook_types,strategy_id=0) | Q(strategy__in=strategys) queryset = queryset.filter(q) @@ -310,8 +310,10 @@ def get(self, request): queryset = queryset.filter(q) level_summary = queryset.values('level').order_by('level').annotate(total=Count('level')) - type_summary = queryset.values('hook_type_id','strategy_id').order_by( - 'hook_type_id').annotate(total=Count('hook_type_id')) + type_summary = queryset.values( + 'hook_type_id', 'strategy_id', 'hook_type__name', + 'strategy__vul_name').order_by('hook_type_id').annotate( + total=Count('hook_type_id')) end['data']['language'] = self.get_languages(queryset.values('agent_id')) @@ -350,10 +352,15 @@ def get_level_name(id): def get_hook_type_name(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 + #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, [strategy_name, hook_type_name])) type_ = list( - filter(lambda x: x is not None, [strategy_name, hook_type_name])) + filter(lambda x: x is not None, [ + obj.get('strategy__vul_name', None), + obj.get('hook_type__name', None) + ])) return type_[0] if type_ else ''