Feat/new metric handler

dongtai_common/migrations/0036_auto_20231013_1051.py

@@ -0,0 +1,34 @@
+# Generated by Django 3.2.20 on 2023-10-13 10:51
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("dongtai_common", "0035_alter_user_phone"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="iastagent",
+            name="jvm_user_dir",
+            field=models.CharField(default="", max_length=1024),
+        ),
+        migrations.CreateModel(
+            name="IastAgentDiskList",
+            fields=[
+                ("id", models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("create_at", models.DateTimeField(auto_now_add=True, null=True)),
+                ("data", models.JSONField(default=dict)),
+                (
+                    "agent",
+                    models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="dongtai_common.iastagent"),
+                ),
+            ],
+            options={
+                "db_table": "iast_agent_disk_list",
+                "managed": True,
+            },
+        ),
+    ]

dongtai_common/models/agent.py

@@ -50,6 +50,7 @@ class IastAgent(models.Model):
     events = models.JSONField(default=get_events)
     department = models.ForeignKey(Department, models.DO_NOTHING)
     allow_report = models.IntegerField(default=1)
+    jvm_user_dir = models.CharField(max_length=1024, default="")
 
     class Meta:
         managed = get_managed()
@@ -91,6 +92,11 @@ class Meta:
         db_table = "iast_agent_event"
 
 
-# class IastAgent(models.Model):
-#
-#    class Meta:
+class IastAgentDiskList(models.Model):
+    agent = models.ForeignKey(IastAgent, on_delete=models.CASCADE)
+    create_at = models.DateTimeField(blank=True, null=True, auto_now_add=True)
+    data = models.JSONField(default=dict)
+
+    class Meta:
+        managed = get_managed()
+        db_table = "iast_agent_disk_list"

dongtai_common/utils/const.py

@@ -3,6 +3,7 @@
 
 # report
 REPORT_HEART_BEAT = 0x01
+REPORT_METRIC = 0x02
 REPORT_SCA = 0x11
 REPORT_VULN_NORNAL = 0x21
 REPORT_VULN_DYNAMIC = 0x22

dongtai_protocol/report/handler/metric_handler.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+# datetime:2020/10/23 11:56
+import logging
+import os
+import time
+
+from dongtai_common.models.agent import (
+    IastAgent,
+    IastAgentDiskList,
+)
+from dongtai_common.utils import const
+from dongtai_protocol.report.handler.heartbeat_handler import update_heartbeat
+from dongtai_protocol.report.handler.report_handler_interface import IReportHandler
+from dongtai_protocol.report.report_handler_factory import ReportHandler
+
+logger = logging.getLogger("dongtai.openapi")
+
+
+@ReportHandler.register(const.REPORT_METRIC)
+class ReportMetricHandler(IReportHandler):
+    def __init__(self):
+        super().__init__()
+        self.req_count = None
+        self.cpu = None
+        self.memory = None
+        self.network = None
+        self.report_queue = None
+        self.method_queue = None
+        self.replay_queue = None
+        self.return_queue = None
+
+    def parse(self):
+        self.cpu = self.detail.get("cpu")
+        self.memory = self.detail.get("memory")
+        self.disk_list = self.detail.get("diskList")
+
+    def has_permission(self):
+        self.agent = IastAgent.objects.filter(id=self.agent_id).first()
+        return self.agent
+
+    def save_heartbeat(self):
+        default_dict = {"dt": int(time.time())}
+        default_dict["memory"] = self.memory
+        default_dict["cpu"] = self.cpu
+        default_dict["disk"] = get_disk_from_disk_list(self.disk_list, self.agent.jvm_user_dir)
+        update_heartbeat.delay(agent_id=self.agent_id, defaults=default_dict)
+        IastAgentDiskList.objects.create(agent=self.agent_id, data=self.disk_list)
+
+    def get_result(self, msg=None):
+        return []
+
+    def save(self):
+        self.save_heartbeat()
+
+    def get_agent(self, agent_id):
+        return IastAgent.objects.filter(id=agent_id).first()
+
+
+def get_disk_from_disk_list(disk_list: list, agent_jvm_dir: str):
+    total_space_bytes = 0
+    usable_space_bytes = 0
+    for disk in disk_list:
+        for partition in disk["partitionList"]:
+            if is_dir_under_another_dir(agent_jvm_dir, partition["mountPoint"]):
+                return get_disk_rate_dict(partition["totalSpaceBytes"], partition["usableSpaceBytes"])
+            total_space_bytes += partition["totalSpaceBytes"]
+            usable_space_bytes += partition["usableSpaceBytes"]
+    return get_disk_rate_dict(total_space_bytes, usable_space_bytes)
+
+
+def get_disk_rate_dict(total_space_bytes: int, usable_space_bytes: int):
+    return (total_space_bytes - usable_space_bytes) / total_space_bytes
+
+
+def is_dir_under_another_dir(subdir: str, parent_dir: str) -> bool:
+    """Returns True if the subdirectory is under the parent directory, False otherwise.
+
+    Args:
+      subdir: The subdirectory path.
+      parent_dir: The parent directory path.
+
+    Returns:
+      True if the subdirectory is under the parent directory, False otherwise.
+    """
+
+    # Normalize the paths.
+    subdir = os.path.normpath(subdir)
+    parent_dir = os.path.normpath(parent_dir)
+
+    # Check if the subdirectory path starts with the parent directory path.
+    return subdir.startswith(parent_dir)

dongtai_protocol/views/agent_register.py

@@ -43,7 +43,7 @@ class AgentRegisterEndPoint(OpenApiEndPoint):
     description = "引擎注册"
 
     @staticmethod
-    def register_agent(token, version, language, project_name, user, project_version):
+    def register_agent(token, version, language, project_name, user, project_version, jvm_user_dir=""):
         project = IastProject.objects.values("id").filter(name=project_name).first()
         is_audit = AgentRegisterEndPoint.get_is_audit()
         project_id = -1
@@ -76,6 +76,7 @@ def register_agent(token, version, language, project_name, user, project_version
                 language=language,
                 is_audit=is_audit,
                 allow_report=allow_report,
+                jvm_user_dir=jvm_user_dir,
             )
         else:
             IastAgent.objects.filter(pk=agent_id).update(
@@ -252,7 +253,7 @@ def post(self, request: Request):
             version_name = param.get("projectVersion", "V1.0")
             version_name = version_name if version_name else "V1.0"
             template_id = param.get("projectTemplateId", None)
-
+            jvm_user_dir = param.get("jvmUserDirectory", "")
             if template_id is not None:
                 template = IastProjectTemplate.objects.filter(pk=template_id).first()
                 if not template:
@@ -294,6 +295,7 @@ def post(self, request: Request):
                     version=version,
                     project_version=project_version,
                     user=user,
+                    jvm_user_dir=jvm_user_dir,
                 )
             else:
                 agent_id = self.register_agent(
@@ -303,6 +305,7 @@ def post(self, request: Request):
                     version=version,
                     user=user,
                     project_version=None,
+                    jvm_user_dir=jvm_user_dir,
                 )
 
             self.register_server(
@@ -367,6 +370,7 @@ def __register_agent(
         language,
         is_audit,
         allow_report,
+        jvm_user_dir="",
     ):
         if exist_project:
             IastAgent.objects.filter(
@@ -390,6 +394,7 @@ def __register_agent(
             is_audit=is_audit,
             allow_report=allow_report,
             department_id=1,
+            jvm_user_dir=jvm_user_dir,
         )
         return agent.id