From b8f2c98439285177489a9960084650ce536296ff Mon Sep 17 00:00:00 2001 From: st1020 Date: Tue, 1 Aug 2023 14:44:09 +0800 Subject: [PATCH 1/2] feat: add report validated sink config --- dongtai_protocol/views/agent_configv2.py | 28 +++++++++++++++++++----- dongtai_web/common.py | 4 ---- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/dongtai_protocol/views/agent_configv2.py b/dongtai_protocol/views/agent_configv2.py index 65845c032..49b5dddaa 100644 --- a/dongtai_protocol/views/agent_configv2.py +++ b/dongtai_protocol/views/agent_configv2.py @@ -1,10 +1,14 @@ +import json +from typing import Any + from django.utils.translation import gettext_lazy as _ from rest_framework import serializers from rest_framework.serializers import ValidationError from dongtai_common.endpoint import OpenApiEndPoint, R from dongtai_common.models.agent import IastAgent -from dongtai_web.common import get_data_gather_data +from dongtai_common.models.profile import IastProfile +from dongtai_conf.patch import patch_point from dongtai_web.utils import extend_schema_with_envcheck @@ -12,8 +16,19 @@ class _AgentConfigArgsSerializer(serializers.Serializer): agent_id = serializers.IntegerField(required=True, help_text=_("Agent id")) -def get_agent_data_gather_config(agent_id): - return get_data_gather_data() +REPORT_VALIDATED_SINK_KEY = "report_validated_sink" +DEFAULT_REPORT_VALIDATED_SINK = {"report_validated_sink": False} + + +def get_report_validated_sink_profile() -> dict[str, bool]: + profile = IastProfile.objects.filter(key=REPORT_VALIDATED_SINK_KEY).values_list("value", flat=True).first() + if profile is None: + IastProfile( + key=REPORT_VALIDATED_SINK_KEY, + value=json.dumps(DEFAULT_REPORT_VALIDATED_SINK), + ).save() + return DEFAULT_REPORT_VALIDATED_SINK + return json.loads(profile) class AgentConfigAllinOneView(OpenApiEndPoint): @@ -29,12 +44,15 @@ def get(self, request): ser.is_valid(True) except ValidationError as e: return R.failure(data=e.detail) - agent = IastAgent.objects.filter(pk=ser.data["agent_id"]).first() + agent_id: int = ser.data["agent_id"] + agent = IastAgent.objects.filter(pk=agent_id).first() if not agent: return R.failure(msg="No agent found.") - data = get_agent_data_gather_config(ser.data["agent_id"]) + data: dict[str, Any] = {} if agent.bind_project is not None and agent.bind_project.enable_log is not None: data["enable_log"] = agent.bind_project.enable_log if agent.bind_project is not None and agent.bind_project.log_level is not None: data["log_level"] = agent.bind_project.log_level + data[REPORT_VALIDATED_SINK_KEY] = get_report_validated_sink_profile()[REPORT_VALIDATED_SINK_KEY] + data, agent_id = patch_point(data, agent_id) return R.success(data=data) diff --git a/dongtai_web/common.py b/dongtai_web/common.py index 34f5003bc..66ac484d3 100644 --- a/dongtai_web/common.py +++ b/dongtai_web/common.py @@ -17,7 +17,3 @@ def get_json_from_iast_profile(key: str, _serializer: type[serializers.Serialize ser = _serializer(data=profile_data) ser.is_valid() return dict(ser.data) - - -def get_data_gather_data() -> dict: - return {} From a86fc739b907740f9f9dfd90cbb217000b78bfe3 Mon Sep 17 00:00:00 2001 From: st1020 Date: Tue, 1 Aug 2023 14:53:51 +0800 Subject: [PATCH 2/2] test: update tests --- test/apiserver/test_agent_config.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/apiserver/test_agent_config.py b/test/apiserver/test_agent_config.py index 62fdbfc27..1e47779c3 100644 --- a/test/apiserver/test_agent_config.py +++ b/test/apiserver/test_agent_config.py @@ -32,11 +32,12 @@ def test_rep_agent_config(self): self.assertEqual(res.status_code, 200) self.assertEqual( data["data"], - {}, + {"report_validated_sink": False}, ) def test_rep_agent_config2(self): agent = IastAgent.objects.filter(pk=self.agent_id).first() + assert agent is not None agent.bind_project.log_level = "INFO" agent.bind_project.enable_log = True agent.bind_project.save() @@ -48,11 +49,12 @@ def test_rep_agent_config2(self): self.assertEqual(res.status_code, 200) self.assertEqual( data["data"], - {"enable_log": True, "log_level": "INFO"}, + {"enable_log": True, "log_level": "INFO", "report_validated_sink": False}, ) def test_rep_agent_config3(self): agent = IastAgent.objects.filter(pk=self.agent_id).first() + assert agent is not None agent.bind_project.log_level = "INFO" agent.bind_project.save() res = self.client.get( @@ -63,5 +65,5 @@ def test_rep_agent_config3(self): self.assertEqual(res.status_code, 200) self.assertEqual( data["data"], - {"log_level": "INFO"}, + {"log_level": "INFO", "report_validated_sink": False}, )