diff --git a/Pipfile b/Pipfile index dd75e78d..30823e75 100644 --- a/Pipfile +++ b/Pipfile @@ -96,6 +96,7 @@ more-itertools = "*" pyotp = "~=2.9.0" qrcode = "~=7.4" django-auth-ldap = "==4.6.0" +tqdm = "*" [dev-packages] diff --git a/Pipfile.lock b/Pipfile.lock index 9b381e0d..e7831b3c 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "f4caca89775b3056c863e1088babb5e9549908d2bcb9bacaae09b826d5e9e54f" + "sha256": "e765376efbd5d0a6e8a22c384e7e008024af5a7d8eb0768db286b7524078e15e" }, "pipfile-spec": 6, "requires": { @@ -1951,6 +1951,14 @@ "markers": "python_version >= '3.8'", "version": "==6.3.3" }, + "tqdm": { + "hashes": [ + "sha256:d302b3c5b53d47bce91fea46679d9c3c6508cf6332229aa1e7d8653723793386", + "sha256:d88e651f9db8d8551a62556d3cff9e3034274ca5d66e93197cf2490e2dcb69c7" + ], + "index": "pypi", + "version": "==4.66.1" + }, "types-awscrt": { "hashes": [ "sha256:477a14565909312fe1de70d0b301548e83c038f436b8a1d7c83729e87cdd0b85", diff --git a/deploy/commands/management/commands/load_hook_strategy.py b/deploy/commands/management/commands/load_hook_strategy.py index 3acf6c4f..2f672f72 100644 --- a/deploy/commands/management/commands/load_hook_strategy.py +++ b/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 .")) diff --git a/dongtai_web/vul_log/vul_log_view.py b/dongtai_web/vul_log/vul_log_view.py index 51d25242..6d9ab108 100644 --- a/dongtai_web/vul_log/vul_log_view.py +++ b/dongtai_web/vul_log/vul_log_view.py @@ -3,6 +3,7 @@ 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 @@ -10,8 +11,8 @@ 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)) diff --git a/requirements.txt b/requirements.txt index 680ac925..5ed92cea 100644 --- a/requirements.txt +++ b/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