feat: new sca select query change.

dongtai_common/endpoint/__init__.py

@@ -28,6 +28,9 @@
 from django.utils.translation import gettext_lazy as _
 from django.db.models import Q, Count
 from typing import Tuple, Dict, Union, TYPE_CHECKING
+from dongtai_common.models.department import Department
+from functools import reduce
+from operator import ior
 
 if TYPE_CHECKING:
     from django.core.paginator import _SupportsPagination
@@ -225,7 +228,12 @@ def get_auth_agents(users):
         :param users:
         :return:
         """
-        return IastAgent.objects.filter(user__in=users)
+        qs = Department.objects.none()
+        users = []
+        qss = [user.get_relative_department() for user in users]
+        departments = reduce(ior, qss, qs)
+        return IastAgent.objects.filter(
+            bind_project__department__in=departments)
         # if isinstance(users, QuerySet):
         #    return IastAgent.objects.filter(user__in=users)
         # else:
@@ -238,7 +246,11 @@ def get_auth_assets(users):
         :param users:
         :return:
         """
-        return Asset.objects.filter(user__in=users, is_del=0)
+        qs = Department.objects.none()
+        users = []
+        qss = [user.get_relative_department() for user in users]
+        departments = reduce(ior,qss, qs)
+        return Asset.objects.filter(department__in=departments, is_del=0)
 
     @staticmethod
     def get_auth_asset_aggrs(auth_assets):
@@ -275,14 +287,14 @@ def get_auth_asset_vuls(assets):
 
     @staticmethod
     def get_auth_and_anonymous_agents(user):
-        query_user = []
-        if user.is_active:
-            query_user = user
-
-        if query_user == []:
-            dt_range_user = User.objects.filter(username=const.USER_BUGENV).first()
-            if dt_range_user:
-                query_user = dt_range_user
+        #        query_user = []
+        #        if user.is_active:
+        #            query_user = user
+        #
+        #        if query_user == []:
+        #            dt_range_user = User.objects.filter(username=const.USER_BUGENV).first()
+        #            if dt_range_user:
+        #                query_user = dt_range_user
         return EndPoint.get_auth_agents_with_user(query_user)
 
 

dongtai_web/aggr_vul/aggr_vul_list.py

@@ -174,10 +174,11 @@ def post(self, request):
 
         except ValidationError as e:
             return R.failure(data=e.detail)
-        user_auth_info = auth_user_list_str(user=request.user,
-                                            user_table="asset")
+        departments = list(request.user.get_relative_department())
+        department_filter_sql = " and {}.department_id in ({})".format(
+            "asset", ",".join(map(lambda x: str(x.id), departments)))
         query_condition = query_condition + \
-            user_auth_info.get("user_condition_str")
+            department_filter_sql
 
         if keywords:
             query_base = "SELECT DISTINCT(vul.id),vul.id,vul.level_id,vul.update_time_desc,vul.update_time," \
@@ -273,7 +274,7 @@ def post(self, request):
             afdistset = Asset.objects.filter(
                 iastvulassetrelation__asset_vul_id__in=vul_ids,
                 iastvulassetrelation__is_del=0,
-                user_id__in=user_auth_info['user_list'],
+                department__in=departments,
                 project_id__gt=0).values(
                     'project_id',
                     'iastvulassetrelation__asset_vul_id').annotate(

dongtai_web/aggr_vul/aggr_vul_summary.py

@@ -12,7 +12,7 @@
 from django.db import connection
 from dongtai_common.common.utils import cached_decorator
 from dongtai_common.models import APP_LEVEL_RISK
-
+from dongtai_common.models.user import User
 
 def get_annotate_sca_common_data(user_id: int, pro_condition: str):
     return get_annotate_sca_base_data(user_id, pro_condition)
@@ -49,9 +49,13 @@ def get_annotate_sca_base_data(user_id: int, pro_condition: str):
         "project": []
     }
     # auth_condition = getAuthBaseQuery(user_id=user_id, table_str="asset")
-    user_auth_info = auth_user_list_str(user_id=user_id, user_table="asset")
+    #user_auth_info = auth_user_list_str(user_id=user_id, user_table="asset")
+    user = User.objects.get(pk=user_id)
+    departments = list(user.get_relative_department())
+    department_filter_sql = " and {}.department_id in ({})".format(
+        "asset", ",".join(map(lambda x: str(x.id), departments)))
     query_condition = " where rel.is_del=0 and asset.project_id>0 " + \
-        user_auth_info.get("user_condition_str") + pro_condition
+        department_filter_sql + pro_condition
     base_join = "left JOIN iast_asset_vul_relation as rel on rel.asset_vul_id=vul.id  " \
                 "left JOIN iast_asset as asset on rel.asset_id=asset.id "
     # level_join = "left JOIN iast_vul_level as level on level.id=vul.level_id "
@@ -172,10 +176,11 @@ def get_annotate_data_es(
     from dongtai_conf import settings
     from dongtai_web.utils import dict_transfrom
     user_id_list = [user_id]
-    auth_user_info = auth_user_list_str(user_id=user_id)
-    user_id_list = auth_user_info['user_list']
+    user = User.objects.get(pk=user_id)
+    departments = list(user.get_relative_department())
+    department_ids = [i.id for i in departments]
     must_query = [
-        Q('terms', asset_user_id=user_id_list),
+        Q('terms', asset_department_id=department_ids),
         Q('terms', asset_vul_relation_is_del=[0]),
         Q('range', asset_project_id={'gt': 0}),
     ]

dongtai_web/dongtai_sca/common/sca_vul.py

@@ -13,7 +13,7 @@ def get_ref(refs) -> list:
     return refs
 
 # 通过asset_vul获取 组件详情信息
-def GetScaVulData(asset_vul, asset_queryset):
+def GetScaVulData(asset_vul, asset_queryset=None):
     data = {'base_info': dict(), 'poc_info': dict()}
     vul_id = asset_vul.id
 

dongtai_web/dongtai_sca/views/package_vul.py

@@ -199,19 +199,17 @@ def get(self, request, vul_id):
         # 组件漏洞基础 数据读取
         asset_vul = IastAssetVul.objects.filter(id=vul_id).first()
         # 用户鉴权
-        auth_users = self.get_auth_users(request.user)
-        asset_queryset = self.get_auth_assets(auth_users)
-
+        departments = request.user.get_relative_department()
         # 判断是否有权限
-        if not asset_vul or not permission_to_read_asset_vul(auth_users, vul_id):
+        if not asset_vul or not permission_to_read_asset_vul(departments, vul_id):
             return R.failure(
                 msg=_('Vul do not exist or no permission to access'))
 
-        data = GetScaVulData(asset_vul, asset_queryset)
+        data = GetScaVulData(asset_vul)
 
         return R.success(data=data)
 
 
-def permission_to_read_asset_vul(users, asset_vul_id: int):
+def permission_to_read_asset_vul(departments, asset_vul_id: int):
     return IastVulAssetRelation.objects.filter(
-        asset__user__in=users, asset_vul_id=asset_vul_id).exists()
+        asset__department__in=departments, asset_vul_id=asset_vul_id).exists()

dongtai_web/views/api_route_cover_rate.py

@@ -55,12 +55,16 @@ def get(self, request):
                 project_id)
         else:
             current_project_version = get_project_version_by_id(version_id)
+        departments = request.user.get_relative_department()
+        projectexist = IastProject.objects.filter(department__in=departments,
+                                                  pk=project_id).first()
+        if not projectexist:
+            return R.failure(_("Parameter error"))
         agents = IastAgent.objects.filter(
-            user__in=auth_users,
             bind_project_id=project_id,
             project_version_id=current_project_version.get("version_id",
                                                            0)).values("id")
-        q = Q(agent__in=agents)
+        q = Q(project_version_id=version_id, project_id=project_id)
         queryset = IastApiRoute.objects.filter(q)
         total = queryset.values("path").distinct().count()
         cover_count = checkcover_batch(queryset, agents)

dongtai_web/views/api_route_search.py

@@ -33,6 +33,7 @@
 from dongtai_web.utils import extend_schema_with_envcheck, get_response_serializer
 import logging
 from dongtai_common.models.strategy import IastStrategyModel
+from dongtai_common.models.project import IastProject
 
 logger = logging.getLogger('dongtai-webapi')
 
@@ -180,6 +181,16 @@ def post(self, request):
             project_version_id=current_project_version.get("version_id",
                                                            0)).values("id")
         q = Q(agent_id__in=[_['id'] for _ in agents])
+        departments = request.user.get_relative_department()
+        projectexist = IastProject.objects.filter(department__in=departments,
+                                                  pk=project_id).first()
+        if not projectexist:
+            return R.failure(_("Parameter error"))
+        agents = IastAgent.objects.filter(
+            bind_project_id=project_id,
+            project_version_id=current_project_version.get("version_id",
+                                                           0)).values("id")
+        q = Q(project_version_id=version_id, project_id=project_id)
         q = q & Q(
             method_id__in=[_['id']
                            for _ in api_methods]) if api_methods != [] else q