Merge pull request #1561 from st1020/feat/remove-some-vul-status

feat: remove some vul status

dongtai_engine/signals/handlers/vul_handler.py

@@ -430,7 +430,7 @@ def handler_replay_vul(vul_meta, vul_level, strategy_id, vul_stack, top_stack,
                        bottom_stack, **kwargs):
     timestamp = int(time.time())
     vul = IastVulnerabilityModel.objects.filter(
-        Q(pk=kwargs['relation_id']) & ~Q(status_id=(3, 4, 5, 6))).first()
+        Q(pk=kwargs['relation_id']) & ~Q(status_id=(3, 5, 6))).first()
     logger.info(
         f'handle vul replay, current strategy:{vul.strategy_id}, target hook_type:{strategy_id}'
     )

dongtai_engine/tasks.py

@@ -10,6 +10,7 @@
 import json
 import time
 from json import JSONDecodeError
+from itertools import groupby
 
 from celery import shared_task
 from celery.apps.worker import logger
@@ -38,6 +39,7 @@
 from dongtai_engine.task_base import replay_payload_data
 from dongtai_engine.common.queryset import get_scan_id, load_sink_strategy, get_agent
 from dongtai_engine.plugins.project_time_update import project_time_stamp_update
+from dongtai_web.vul_log.vul_log import log_recheck_vul
 
 RETRY_INTERVALS = [10, 30, 90]
 
@@ -145,9 +147,12 @@ def search_and_save_vul(engine: Optional[VulEngine],
             replay_id = method_pool_model.replay_id
             relation_id = method_pool_model.relation_id
             timestamp = int(time.time())
-            IastVulnerabilityModel.objects.filter(id=relation_id).update(
-                status_id=settings.IGNORE,
-                latest_time=timestamp
+            vul = IastVulnerabilityModel.objects.filter(id=relation_id).get()
+            log_recheck_vul(
+                vul.agent.user.id,
+                vul.agent.user.username,
+                [vul.id],
+                '已忽略',
             )
             IastReplayQueue.objects.filter(id=replay_id).update(
                 state=const.SOLVED,
@@ -369,7 +374,7 @@ def update_agent_status():
     """
     from dongtai_engine.plugins.engine_status_change import after_agent_status_update, before_agent_status_update
     before_agent_status_update()
-    logger.info(f'检测引擎状态更新开始')
+    logger.info('检测引擎状态更新开始')
     timestamp = int(time.time())
     running_agents_ids = list(
         IastAgent.objects.values("id").filter(online=1).values_list(
@@ -385,11 +390,19 @@ def update_agent_status():
         update_time__lte=timestamp - 60 * 5,
         verify_time__isnull=True,
         replay_type=1).values('relation_id').distinct()
-    IastVulnerabilityModel.objects.filter(
-        Q(pk__in=vul_id_qs)
-        & ~Q(status_id__in=(3, 4, 5, 6))).update(status_id=7)
+    vuls = IastVulnerabilityModel.objects.filter(
+        Q(pk__in=vul_id_qs) & ~Q(status_id__in=(3, 5, 6))
+    ).select_related("agent__user")
+    for _, vul_list in groupby(vuls, lambda x: x.agent.user_id):
+        vul_list = list(vul_list)
+        log_recheck_vul(
+            vul_list[0].agent.user.id,
+            vul_list[0].agent.user.username,
+            list(map(lambda x: x.id, vul_list)),
+            "验证失败",
+        )
     logger.info("update offline agent: %s", stop_agent_ids)
-    logger.info(f'检测引擎状态更新成功')
+    logger.info('检测引擎状态更新成功')
     after_agent_status_update()
 
 @shared_task(queue='dongtai-periodic-task')

dongtai_web/views/vul_recheck_v2.py

@@ -10,10 +10,8 @@
 from dongtai_common.models.project import IastProject
 from dongtai_common.models.replay_queue import IastReplayQueue
 from dongtai_common.models.vulnerablity import IastVulnerabilityModel
-from dongtai_common.utils.validate import Validate
 from dongtai_web.aggregation.aggregation_common import turnIntListOfStr
 from dongtai_web.utils import extend_schema_with_envcheck, get_response_serializer
-from rest_framework import serializers
 
 from dongtai_common.endpoint import R
 from dongtai_common.utils import const
@@ -158,7 +156,7 @@ def post(self, request):
 
             queryset = IastVulnerabilityModel.objects.filter(
                 Q(is_del=0, project__department__in=department)
-                & ~Q(status_id__in=(3, 4, 5, 6)))
+                & ~Q(status_id__in=(3, 5, 6)))
             ids_list = turnIntListOfStr(vul_ids)
 
             vul_queryset = queryset.filter(id__in=ids_list)

test/core/tasks.py

@@ -1,7 +1,10 @@
+import json
 import unittest
+from rest_framework.test import APITestCase
 
 from test import DongTaiTestCase
-
+from test.apiserver.test_agent_base import REGISTER_JSON, METHODPOOL_JSON, gzipdata
+from dongtai_common.models.user import User
 
 class MyTestCase(DongTaiTestCase):
     def test_something(self):
@@ -105,5 +108,33 @@ def parse_path(self):
         print(project_cookies)
 
 
+class MultiUserTestCase(APITestCase):
+    def test_update_agent_status_multi_user(self):
+        from dongtai_engine.tasks import update_agent_status
+        user1 = User.objects.filter(pk=1).get()
+        assert user1 is not None
+        self.client.force_authenticate(user=user1)
+        data1 = self.register_agent(name='test1')
+
+        User(id=2, username="test", phone="123456789").save()
+        user2 = User.objects.filter(pk=2).get()
+        assert user2 is not None
+        self.client.force_authenticate(user=user2)
+        data2 = self.register_agent(name='test2')
+
+        update_agent_status()
+
+    def register_agent(self, **kwargs):
+        register_data = REGISTER_JSON
+        register_data.update(kwargs)
+        data = gzipdata(REGISTER_JSON)
+        response = self.client.post('http://testserver/api/v1/agent/register',
+                                    data=data,
+                                    HTTP_CONTENT_ENCODING='gzip',
+                                    content_type='application/json',
+                                    )
+        return json.loads(response.content)['data']
+
+
 if __name__ == '__main__':
     unittest.main()