feat load vul fix

Pipfile

@@ -96,6 +96,7 @@ more-itertools = "*"
 pyotp = "~=2.9.0"
 qrcode = "~=7.4"
 django-auth-ldap = "==4.6.0"
+tqdm = "*"
 
 [dev-packages]
 

deploy/commands/management/commands/load_hook_strategy.py

@@ -4,6 +4,7 @@
 
 from django.core.management.base import BaseCommand
 from django.db.models import Q
+from tqdm import tqdm
 
 from dongtai_common.models.hook_strategy import HookStrategy
 from dongtai_common.models.hook_type import HookType
@@ -29,7 +30,7 @@ def handle(self, *args, **options):
             with open(os.path.join(POLICY_DIR, "sensitive_info_strategy.json")) as fp:
                 full_strategies.extend(json.load(fp, object_pairs_hook=OrderedDict))
         strategy_dict = {}
-        for strategy in full_strategies:
+        for strategy in tqdm(full_strategies, desc="strategy"):
             if IastStrategyModel.objects.filter(
                 vul_type=strategy["vul_type"],
                 system_type=1,
@@ -85,7 +86,7 @@ def handle(self, *args, **options):
             with open(os.path.join(POLICY_DIR, f"{k.lower()}_hooktype.json")) as fp:
                 hooktypes = json.load(fp, object_pairs_hook=OrderedDict)
             hooktype_dict = {}
-            for hook_type in hooktypes:
+            for hook_type in tqdm(hooktypes, desc="hook_type"):
                 if HookType.objects.filter(
                     value=hook_type["value"],
                     type=hook_type["type"],
@@ -119,7 +120,7 @@ def handle(self, *args, **options):
             HookStrategy.objects.filter(language_id=v, system_type=1).update(system_type=0)
             with open(os.path.join(POLICY_DIR, f"{k.lower()}_full_policy.json")) as fp:
                 full_policy = json.load(fp, object_pairs_hook=OrderedDict)
-            for policy in full_policy:
+            for policy in tqdm(full_policy, desc="policy"):
                 if policy["type"] == 4:
                     if policy["value"] not in strategy_dict:
                         continue
@@ -189,4 +190,17 @@ def handle(self, *args, **options):
                 )
                 sensitive_info_rule_ids.append(obj.pk)
         IastSensitiveInfoRule.objects.filter(~Q(id__in=sensitive_info_rule_ids), system_type=1).delete()
+        update_fix_vul_dict = {}
+        if os.path.exists(os.path.join(POLICY_DIR, "vul_fix_extend.json")):
+            with open(os.path.join(POLICY_DIR, "vul_fix_extend.json")) as fp:
+                update_fix_vul_dict = json.load(fp, object_pairs_hook=OrderedDict)
+        need_update_fix_vul_strategies = IastStrategyModel.objects.filter(
+            vul_type__in=list(update_fix_vul_dict.keys()), system_type=1
+        ).all()
+        for strategy in tqdm(need_update_fix_vul_strategies, desc="vul_fix"):
+            if strategy.vul_type in update_fix_vul_dict:
+                update_content = update_fix_vul_dict[strategy.vul_type]
+                strategy.vul_fix = update_content["vul_fix"]
+                strategy.vul_fix_zh = update_content["vul_fix"]
+                strategy.vul_fix_en = update_content["vul_fix"]
         self.stdout.write(self.style.SUCCESS("Successfully load strategy ."))

dongtai_web/vul_log/vul_log_view.py

@@ -3,15 +3,16 @@
 from django.utils.translation import gettext_lazy as _
 from drf_spectacular.utils import extend_schema
 from rest_framework import serializers, viewsets
+from rest_framework.serializers import ValidationError
 
 from dongtai_common.endpoint import R, UserEndPoint
 from dongtai_common.models.iast_vul_log import IastVulLog
 from dongtai_web.common import VulType
 
 
 class VulLogListArgsSerializer(serializers.Serializer):
-    vul_type = serializers.IntegerField(help_text="漏洞类型")
-    msg_type = serializers.IntegerField(required=False, help_text="消息类型")
+    vul_type = serializers.IntegerField(min_value=1, max_value=2, help_text="漏洞类型")
+    msg_type = serializers.IntegerField(min_value=1, max_value=5, required=False, help_text="消息类型")
 
 
 class VulLogViewSet(UserEndPoint, viewsets.ViewSet):
@@ -25,6 +26,12 @@ class VulLogViewSet(UserEndPoint, viewsets.ViewSet):
     )
     def list(self, request, vul_id):
         data = []
+        ser = VulLogListArgsSerializer(data=request.GET)
+        try:
+            if ser.is_valid(True):
+                pass
+        except ValidationError as e:
+            return R.failure(data=e.detail)
         auth_users = self.get_auth_users(request.user)
         vul_type = VulType(int(request.query_params.get("vul_type", 1)))
         msg_type = int(request.query_params.get("msg_type", 1))

requirements.txt

@@ -135,6 +135,7 @@ sqlparse==0.4.4 ; python_version >= '3.5'
 tablib[html,ods,xls,xlsx,yaml]==3.5.0 ; python_version >= '3.8'
 tomli==2.0.1 ; python_version < '3.11'
 tornado==6.3.3 ; python_version >= '3.8'
+tqdm==4.66.1
 types-awscrt==0.19.2 ; python_version >= '3.7' and python_version < '4.0'
 types-pymysql==1.1.0.1
 types-pyopenssl==23.2.0.2