Merge pull request #595 from Bidaya0/feature/agent_config_v2_final

Feature/agent config v2 final

apiserver/views/agent_config.py

@@ -42,7 +42,7 @@ def post(self, request):
                     cluster_name__in=('', server.cluster_name),
                     cluster_version__in=('', server.cluster_version),
                     hostname__in=('', server.hostname),
-                    ip__in=('', server.ip)).order_by('-priority').first()
+                    ip__in=('', server.ip)).order_by('priority').first()
                 if config:
                     data = config.details
 
@@ -64,7 +64,6 @@ class AgentConfigv2View(OpenApiEndPoint):
 
     def post(self, request):
         try:
-            # agent_id = request.data.get('agentId', None)
             param = parse_data(request.read())
             agent_id = int(param.get('agentId', None))
             if agent_id is None:
@@ -76,7 +75,7 @@ def post(self, request):
             return R.success(msg=_('Successfully'), data={})
         res = get_agent_config(agent_id)
         if isinstance(res, Err):
-            return R.success(msg=_(Err.value), data={})
+            return R.success(msg=_(res.value), data={})
         agent_config = res.value
         return R.success(msg=_('Successfully'), data=agent_config)
 
@@ -102,7 +101,8 @@ def get_agent_filter_details(agent_id):
 def get_agent_config_by_scan(agent_id: int, mg: MetricGroup) -> Result:
     agent_detail = get_agent_filter_details(agent_id)
     queryset = IastCircuitConfig.objects.filter(
-        is_deleted=0, metric_group=mg).order_by('-priority').only('id')
+        is_deleted=0, metric_group=mg,
+        is_enable=1).order_by('-priority').only('id')
     for i in queryset:
         result_list = []
         for target in IastCircuitTarget.objects.filter(
@@ -149,6 +149,10 @@ def get_agent_config(agent_id: int) -> Result:
         data[mg.name.lower() +
              "IsUninstall"] = True if config.deal == DealType.UNLOAD else False
         interval_list.append(config.interval)
+    # if interval_list is [], there is mean no config found here. 
+    # because interval is required in create config.
+    if not interval_list:
+        return Err('No config found')
     data["performanceLimitRiskMaxMetricsCount"] = min(interval_list)
     return Ok(data)
 

dongtai/models/agent_config.py

@@ -62,13 +62,16 @@ class MetricType(IntegerChoices):
     daemonThreadCount = 7, _("守护线程数阈值")
     dongTaiThreadCount = 8, _("洞态IAST线程数阈值")
     hookLimitTokenPerSecond = 9, _("单请求HOOK限流")
-    heavyTrafficLimitTokenPerSecond = 10, _("高频HOOK限流")
+    heavyTrafficLimitTokenPerSecond = 10, _("每秒限制处理请求数量（QPS）")
 
 class IastCircuitConfig(models.Model):
     user = models.ForeignKey(User, models.DO_NOTHING)
     name = models.CharField(max_length=200, blank=True, null=True)
     metric_types = models.CharField(max_length=2000, blank=True, null=True)
-    targets = models.CharField(max_length=2000, blank=True, null=True)
+    target_types = models.CharField(max_length=2000,
+                                    blank=True,
+                                    null=True,
+                                    db_column='targets')
     system_type = models.IntegerField(blank=True, null=True)
     is_enable = models.IntegerField(blank=True, null=True)
     is_deleted = models.IntegerField(default=0, blank=True, null=True)

iast/threshold/config_setting.py

@@ -104,6 +104,7 @@ def post(self, request):
 )
 from collections.abc import Iterable
 from inflection import underscore
+from functools import partial
 
 def intable_validate(value):
     try:
@@ -132,16 +133,21 @@ class AgentConfigSettingV2Serializer(serializers.Serializer):
         child=AgentConfigSettingV2MetricSerializer())
     interval = serializers.IntegerField()
     deal = serializers.ChoiceField(DealType.choices)
-    is_enable = serializers.ChoiceField(DealType.choices)
+    is_enable = serializers.IntegerField()
 
 
-from django.db.models import Max
+from django.db.models import Max, Min
+from django.forms.models import model_to_dict
 
 
-def get_priority_now() -> int:
+def get_priority_max_now() -> int:
     res = IastCircuitConfig.objects.all().aggregate(Max("priority"))
     return res["priority__max"] + 1
 
+def get_priority_min_now() -> int:
+    res = IastCircuitConfig.objects.all().aggregate(Min("priority"))
+    return res["priority__min"] - 1
+
 def config_create(data, user):
     fields = ('name', 'metric_group', 'is_enable', 'deal',
               "interval")
@@ -150,7 +156,8 @@ def config_create(data, user):
     targets = get_targets(data['targets'])
     obj = IastCircuitConfig.objects.create(**filted_data,
                                            metric_types=metric_types,
-                                           targets=targets,
+                                           target_types=targets,
+                                           priority=get_priority_max_now(),
                                            user=user)
     for i in data['targets']:
         create_target(i, obj)
@@ -168,7 +175,7 @@ def config_update(data, config_id):
     IastCircuitConfig.objects.filter(
         pk=config_id).update(**filted_data,
                              metric_types=metric_types,
-                             targets=targets)
+                             target_types=targets)
     IastCircuitTarget.objects.filter(
             circuit_config_id=config_id).delete()
     IastCircuitMetric.objects.filter(
@@ -189,8 +196,6 @@ def create_target(target: dict, circuit_config: IastCircuitConfig):
     IastCircuitTarget.objects.create(circuit_config=circuit_config, **target)
 
 
-
-
 def get_metric_types(metrics):
     str_list = []
     for metric in metrics:
@@ -202,12 +207,63 @@ def get_targets(targets):
     str_list = []
     for target in targets:
         str_list.append(str(TargetType(target['target_type']).label))
-    return str(_("、")).join(str_list)
+    res = str(_("、")).join(str_list)
+    if not res:
+        return str(_("全部"))
+    return res
 
 
 def get_data_from_dict_by_key(dic: dict, fields: Iterable) -> dict:
     return {i: dic[i] for i in fields}
 
+from django.db.models import F
+
+
+def set_config_top(config_id):
+    config = IastCircuitConfig.objects.filter(pk=config_id).first()
+    IastCircuitConfig.objects.filter(priority__lt=config.priority).update(
+        priority=F('priority') + 1)
+    config.priority = get_priority_min_now()
+    config.save()
+
+
+#when target_priority < config.priorty
+def set_config_change_lt(config_id, target_priority: int):
+    config = IastCircuitConfig.objects.filter(pk=config_id).first()
+    IastCircuitConfig.objects.filter(
+        priority__gte=target_priority,
+        priority__lt=config.priority).update(priority=F('priority') + 1)
+    config.priority = target_priority
+    config.save()
+
+
+def set_config_top(config_id):
+    return set_config_change_lt(config_id,
+                                target_priority=get_priority_min_now())
+
+
+#when target_priority > config.priorty
+def set_config_change_gt(config_id, target_priority: int):
+    config = IastCircuitConfig.objects.filter(pk=config_id).first()
+    IastCircuitConfig.objects.filter(
+        priority__lte=target_priority,
+        priority__gt=config.priority).update(priority=F('priority') - 1)
+    config.priority = target_priority
+    config.save()
+
+
+def set_config_bottom(config_id):
+    set_config_change_gt(config_id, target_priority=get_priority_max_now())
+
+
+def set_config_change_proprity(config_id, priority_range: list):
+    config = IastCircuitConfig.objects.filter(pk=config_id).first()
+    if min(priority_range) > config.priority:
+        set_config_change_gt(config.id, min(priority_range))
+    if max(priority_range) < config.priority:
+        set_config_change_lt(config.id, max(priority_range))
+from webapi.settings import DEFAULT_CIRCUITCONFIG
+
 
 class AgentThresholdConfigV2(UserEndPoint, viewsets.ViewSet):
     name = "api-v1-agent-threshold-config-setting-v2"
@@ -245,9 +301,16 @@ def list(self, request):
         page = request.query_params.get('page', 1)
         page_size = request.query_params.get("page_size", 10)
         queryset = IastCircuitConfig.objects.filter(
-            is_deleted=0).order_by('-priority').values()
+            is_deleted=0).order_by('priority').prefetch_related(
+                'iastcircuittarget_set', 'iastcircuitmetric_set').all()
         page_summary, page_data = self.get_paginator(queryset, page, page_size)
-        return R.success(page=page_summary, data=list(page_data))
+        obj_list = []
+        for data in page_data:
+            obj = model_to_dict(data)
+            obj['targets'] = list(data.iastcircuittarget_set.values().all())
+            obj['metrics'] = list(data.iastcircuitmetric_set.values().all())
+            obj_list.append(obj)
+        return R.success(page=page_summary, data=obj_list)
 
     def update(self, request, pk):
         ser = AgentConfigSettingV2Serializer(data=request.data)
@@ -260,22 +323,52 @@ def update(self, request, pk):
         return R.success()
 
     def reset(self, request, pk):
-        return R.success()
+        if IastCircuitConfig.objects.filter(pk=pk).exists():
+            config = IastCircuitConfig.objects.filter(pk=pk, ).first()
+            mg = MetricGroup(config.metric_group)
+            data = DEFAULT_CIRCUITCONFIG[mg.name]
+            config_update(data, pk)
+            return R.success()
+        return R.failure
+
+    def change_priority(self, request, pk):
+        type_ = request.data.get('type')
+        priority_range = request.data.get('priority_range')
+        if IastCircuitConfig.objects.filter(pk=pk).exists():
+            if type_ == 1:
+                set_config_top(pk)
+                return R.success()
+            if type_ == 2 and priority_range:
+                set_config_change_proprity(pk, priority_range)
+                return R.success()
+            if type_ == 3:
+                set_config_bottom(pk)
+                return R.success()
+        return R.failure()
 
     def delete(self, request, pk):
         IastCircuitConfig.objects.filter(pk=pk).update(is_deleted=1)
         return R.success()
 
     def enum(self, request, enumname):
-        able_to_search = (MetricType, MetricGroup, TargetOperator,
-                          MetricOperator)
+        able_to_search = (TargetType, MetricType, MetricGroup, TargetOperator,
+                          MetricOperator, DealType)
         able_to_search_dict = {
             underscore(item.__name__): item
             for item in able_to_search
         }
         return R.success(data=convert_choices_to_value_dict(
             able_to_search_dict.get(enumname)))
 
+    def enumall(self, request):
+        able_to_search = (TargetType, MetricType, MetricGroup, TargetOperator,
+                          MetricOperator, DealType)
+        res = {
+            underscore(item.__name__): convert_choices_to_value_dict(item)
+            for item in able_to_search
+        }
+        return R.success(data=res)
+
 
 def convert_choices_to_dict(choices):
     fields = ['value', 'name', 'label']

iast/urls.py

@@ -342,6 +342,18 @@
              "post": "create",
              "get": "list"
          })),
+    path('circuit_config/enum/all',
+         AgentThresholdConfigV2.as_view({
+             "get": "enumall"
+         })),
+    path('circuit_config/<int:pk>/priority',
+         AgentThresholdConfigV2.as_view({
+             "put": "change_priority"
+         })),
+    path('circuit_config/<int:pk>/reset',
+         AgentThresholdConfigV2.as_view({
+             "put": "reset"
+         })),
     path('circuit_config/enum/<str:enumname>',
          AgentThresholdConfigV2.as_view({
              "get": "enum"

test/apiserver/test_config.py

@@ -1,6 +1,7 @@
 from rest_framework.test import APITestCase
 from apiserver.views.agent_config import get_agent_config
 from apiserver.views.agent_config import *
+from dongtai.models.user import User
 
 
 class VulDetailTestCase(APITestCase):
@@ -15,3 +16,9 @@ def test_agent_detail_retrieve(self):
     def test_target_filter(self):
         res = get_agent_config_by_scan(1, 2)
         print(res)
+
+
+    def test_agent_config_request(self):
+        self.user = User.objects.filter(pk=1).first()
+        self.client.force_authenticate(user=self.user)
+        response = self.client.post('/api/v1/agent/thresholdv2')